refactor(handler): extract RowSet formatting from scheduler to handler (#5802)

* lift to_pg_rows

* cleanup `format: bool`

* no enum wrapping local & distributed

* enum stream without boxing for local or distributed

* comments

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: xiangjinwu

Cargo.lock

@@ -28,6 +28,7 @@ parking_lot = "0.12"
 parse-display = "0.6"
 paste = "1"
 pgwire = { path = "../utils/pgwire" }
+pin-project-lite = "0.2"
 prometheus = { version = "0.13", features = ["process"] }
 prost = "0.11"
 rand = "0.8"

src/frontend/Cargo.toml

@@ -25,6 +25,7 @@ use pgwire::types::Row;
 use risingwave_common::error::{ErrorCode, Result};
 use risingwave_sqlparser::ast::{DropStatement, ObjectType, Statement};
 
+use self::util::DataChunkToRowSetAdapter;
 use crate::scheduler::{DistributedQueryStream, LocalQueryStream};
 use crate::session::{OptimizerContext, SessionImpl};
 use crate::utils::WithOptions;
@@ -60,8 +61,8 @@ pub mod variable;
 pub type RwPgResponse = PgResponse<PgResponseStream>;
 
 pub enum PgResponseStream {
-    LocalQuery(LocalQueryStream),
-    DistributedQuery(DistributedQueryStream),
+    LocalQuery(DataChunkToRowSetAdapter<LocalQueryStream>),
+    DistributedQuery(DataChunkToRowSetAdapter<DistributedQueryStream>),
     Rows(BoxStream<'static, RowSetResult>),
 }
 

src/frontend/src/handler/mod.rs

@@ -25,17 +25,16 @@ use risingwave_sqlparser::ast::Statement;
 use super::{PgResponseStream, RwPgResponse};
 use crate::binder::{Binder, BoundSetExpr, BoundStatement};
 use crate::handler::privilege::{check_privileges, resolve_privileges};
-use crate::handler::util::to_pg_field;
+use crate::handler::util::{to_pg_field, DataChunkToRowSetAdapter};
 use crate::planner::Planner;
 use crate::scheduler::plan_fragmenter::Query;
 use crate::scheduler::{
-    BatchPlanFragmenter, ExecutionContext, ExecutionContextRef, LocalQueryExecution,
+    BatchPlanFragmenter, DistributedQueryStream, ExecutionContext, ExecutionContextRef,
+    LocalQueryExecution, LocalQueryStream,
 };
 use crate::session::{OptimizerContext, OptimizerContextRef, SessionImpl};
 use crate::PlanRef;
 
-pub type QueryResultSet = PgResponseStream;
-
 pub fn gen_batch_query_plan(
     session: &SessionImpl,
     context: OptimizerContextRef,
@@ -116,9 +115,17 @@ pub async fn handle_query(
     tracing::trace!("Generated query after plan fragmenter: {:?}", &query);
 
     let mut row_stream = match query_mode {
-        QueryMode::Local => local_execute(session.clone(), query, format).await?,
+        QueryMode::Local => PgResponseStream::LocalQuery(DataChunkToRowSetAdapter::new(
+            local_execute(session.clone(), query).await?,
+            format,
+        )),
         // Local mode do not support cancel tasks.
-        QueryMode::Distributed => distribute_execute(session.clone(), query, format).await?,
+        QueryMode::Distributed => {
+            PgResponseStream::DistributedQuery(DataChunkToRowSetAdapter::new(
+                distribute_execute(session.clone(), query).await?,
+                format,
+            ))
+        }
     };
 
     let rows_count = match stmt_type {
@@ -183,21 +190,16 @@ fn to_statement_type(stmt: &Statement) -> StatementType {
 pub async fn distribute_execute(
     session: Arc<SessionImpl>,
     query: Query,
-    format: bool,
-) -> Result<QueryResultSet> {
+) -> Result<DistributedQueryStream> {
     let execution_context: ExecutionContextRef = ExecutionContext::new(session.clone()).into();
     let query_manager = execution_context.session().env().query_manager().clone();
     query_manager
-        .schedule(execution_context, query, format)
+        .schedule(execution_context, query)
         .await
         .map_err(|err| err.into())
 }
 
-async fn local_execute(
-    session: Arc<SessionImpl>,
-    query: Query,
-    format: bool,
-) -> Result<QueryResultSet> {
+async fn local_execute(session: Arc<SessionImpl>, query: Query) -> Result<LocalQueryStream> {
     let front_env = session.env();
 
     // Acquire hummock snapshot for local execution.
@@ -211,7 +213,7 @@ async fn local_execute(
     // TODO: Passing sql here
     let execution =
         LocalQueryExecution::new(query, front_env.clone(), "", epoch, session.auth_context());
-    let rsp = Ok(execution.stream_rows(format));
+    let rsp = Ok(execution.stream_rows());
 
     // Release hummock snapshot for local execution.
     hummock_snapshot_manager.release(epoch, &query_id).await;

src/frontend/src/handler/query.rs

@@ -12,14 +12,70 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
 use bytes::Bytes;
+use futures::Stream;
 use itertools::Itertools;
 use pgwire::pg_field_descriptor::{PgFieldDescriptor, TypeOid};
+use pgwire::pg_response::RowSetResult;
+use pgwire::pg_server::BoxedError;
 use pgwire::types::Row;
+use pin_project_lite::pin_project;
 use risingwave_common::array::DataChunk;
 use risingwave_common::catalog::{ColumnDesc, Field};
 use risingwave_common::types::{DataType, ScalarRefImpl};
 
+pin_project! {
+    /// Wrapper struct that converts a stream of DataChunk to a stream of RowSet based on formatting
+    /// parameters.
+    ///
+    /// This is essentially `StreamExt::map(self, move |res| res.map(|chunk| to_pg_rows(chunk,
+    /// format)))` but we need a nameable type as part of [`super::PgResponseStream`], but we cannot
+    /// name the type of a closure.
+    pub struct DataChunkToRowSetAdapter<VS>
+    where
+        VS: Stream<Item = Result<DataChunk, BoxedError>>,
+    {
+        #[pin]
+        chunk_stream: VS,
+        format: bool,
+    }
+}
+impl<VS> DataChunkToRowSetAdapter<VS>
+where
+    VS: Stream<Item = Result<DataChunk, BoxedError>>,
+{
+    pub fn new(chunk_stream: VS, format: bool) -> Self {
+        Self {
+            chunk_stream,
+            format,
+        }
+    }
+}
+
+impl<VS> Stream for DataChunkToRowSetAdapter<VS>
+where
+    VS: Stream<Item = Result<DataChunk, BoxedError>>,
+{
+    type Item = RowSetResult;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let mut this = self.project();
+        match this.chunk_stream.as_mut().poll_next(cx) {
+            Poll::Pending => Poll::Pending,
+            Poll::Ready(chunk) => match chunk {
+                Some(chunk_result) => match chunk_result {
+                    Ok(chunk) => Poll::Ready(Some(Ok(to_pg_rows(chunk, *this.format)))),
+                    Err(err) => Poll::Ready(Some(Err(err))),
+                },
+                None => Poll::Ready(None),
+            },
+        }
+    }
+}
+
 /// Format scalars according to postgres convention.
 fn pg_value_format(d: ScalarRefImpl<'_>, format: bool) -> Bytes {
     // format == false means TEXT format
@@ -34,7 +90,7 @@ fn pg_value_format(d: ScalarRefImpl<'_>, format: bool) -> Bytes {
     }
 }
 
-pub fn to_pg_rows(chunk: DataChunk, format: bool) -> Vec<Row> {
+fn to_pg_rows(chunk: DataChunk, format: bool) -> Vec<Row> {
     chunk
         .rows()
         .map(|r| {

src/frontend/src/handler/util.rs

@@ -20,8 +20,7 @@ use std::task::{Context, Poll};
 
 use futures::{Stream, StreamExt};
 use futures_async_stream::try_stream;
-use pgwire::pg_response::RowSetResult;
-use pgwire::pg_server::{Session, SessionId};
+use pgwire::pg_server::{BoxedError, Session, SessionId};
 use risingwave_batch::executor::BoxedDataChunkStream;
 use risingwave_common::array::DataChunk;
 use risingwave_common::error::RwError;
@@ -32,36 +31,23 @@ use tracing::debug;
 
 use super::QueryExecution;
 use crate::catalog::catalog_service::CatalogReader;
-use crate::handler::query::QueryResultSet;
-use crate::handler::util::to_pg_rows;
-use crate::handler::PgResponseStream;
 use crate::scheduler::plan_fragmenter::{Query, QueryId};
 use crate::scheduler::worker_node_manager::WorkerNodeManagerRef;
 use crate::scheduler::{ExecutionContextRef, HummockSnapshotManagerRef, SchedulerResult};
 
 pub struct DistributedQueryStream {
     chunk_rx: tokio::sync::mpsc::Receiver<SchedulerResult<DataChunk>>,
-    format: bool,
-}
-
-impl DistributedQueryStream {
-    pub fn new(
-        chunk_rx: tokio::sync::mpsc::Receiver<SchedulerResult<DataChunk>>,
-        format: bool,
-    ) -> Self {
-        Self { chunk_rx, format }
-    }
 }
 
 impl Stream for DistributedQueryStream {
-    type Item = RowSetResult;
+    type Item = Result<DataChunk, BoxedError>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         match self.chunk_rx.poll_recv(cx) {
             Poll::Pending => Poll::Pending,
             Poll::Ready(chunk) => match chunk {
                 Some(chunk_result) => match chunk_result {
-                    Ok(chunk) => Poll::Ready(Some(Ok(to_pg_rows(chunk, self.format)))),
+                    Ok(chunk) => Poll::Ready(Some(Ok(chunk))),
                     Err(err) => Poll::Ready(Some(Err(Box::new(err)))),
                 },
                 None => Poll::Ready(None),
@@ -117,8 +103,7 @@ impl QueryManager {
         &self,
         context: ExecutionContextRef,
         query: Query,
-        format: bool,
-    ) -> SchedulerResult<QueryResultSet> {
+    ) -> SchedulerResult<DistributedQueryStream> {
         let query_id = query.query_id().clone();
         let epoch = self
             .hummock_snapshot_manager
@@ -157,7 +142,7 @@ impl QueryManager {
 
         // TODO: Clean up queries status when ends. This should be done lazily.
 
-        Ok(query_result_fetcher.stream_from_channel(format))
+        Ok(query_result_fetcher.stream_from_channel())
     }
 
     pub fn cancel_queries_in_session(&self, session_id: SessionId) {
@@ -225,11 +210,10 @@ impl QueryResultFetcher {
         Box::pin(self.run_inner())
     }
 
-    fn stream_from_channel(self, format: bool) -> QueryResultSet {
-        PgResponseStream::DistributedQuery(DistributedQueryStream {
+    fn stream_from_channel(self) -> DistributedQueryStream {
+        DistributedQueryStream {
             chunk_rx: self.chunk_rx,
-            format,
-        })
+        }
     }
 }
 

src/frontend/src/scheduler/distributed/query_manager.rs

@@ -21,7 +21,7 @@ use std::task::{Context, Poll};
 use futures::Stream;
 use futures_async_stream::try_stream;
 use itertools::Itertools;
-use pgwire::pg_response::RowSetResult;
+use pgwire::pg_server::BoxedError;
 use risingwave_batch::executor::{BoxedDataChunkStream, ExecutorBuilder};
 use risingwave_batch::task::TaskId;
 use risingwave_common::array::DataChunk;
@@ -38,8 +38,6 @@ use tracing::debug;
 use uuid::Uuid;
 
 use super::plan_fragmenter::{PartitionInfo, QueryStageRef};
-use crate::handler::query::QueryResultSet;
-use crate::handler::util::to_pg_rows;
 use crate::optimizer::plan_node::PlanNodeType;
 use crate::scheduler::plan_fragmenter::{ExecutionPlanNode, Query, StageId};
 use crate::scheduler::task_context::FrontendBatchTaskContext;
@@ -48,27 +46,17 @@ use crate::session::{AuthContext, FrontendEnv};
 
 pub struct LocalQueryStream {
     data_stream: BoxedDataChunkStream,
-    format: bool,
-}
-
-impl LocalQueryStream {
-    pub fn new(data_stream: BoxedDataChunkStream, format: bool) -> Self {
-        Self {
-            data_stream,
-            format,
-        }
-    }
 }
 
 impl Stream for LocalQueryStream {
-    type Item = RowSetResult;
+    type Item = Result<DataChunk, BoxedError>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         match self.data_stream.as_mut().poll_next(cx) {
             Poll::Pending => Poll::Pending,
             Poll::Ready(chunk) => match chunk {
                 Some(chunk_result) => match chunk_result {
-                    Ok(chunk) => Poll::Ready(Some(Ok(to_pg_rows(chunk, self.format)))),
+                    Ok(chunk) => Poll::Ready(Some(Ok(chunk))),
                     Err(err) => Poll::Ready(Some(Err(Box::new(err)))),
                 },
                 None => Poll::Ready(None),
@@ -134,11 +122,10 @@ impl LocalQueryExecution {
         Box::pin(self.run_inner())
     }
 
-    pub fn stream_rows(self, format: bool) -> QueryResultSet {
-        QueryResultSet::LocalQuery(LocalQueryStream {
+    pub fn stream_rows(self) -> LocalQueryStream {
+        LocalQueryStream {
             data_stream: self.run(),
-            format,
-        })
+        }
     }
 
     /// Convert query to plan fragment.