feat(flow): query table schema&refactor (#3943)

* feat: get table info

* feat: remove new&unwrap

* chore: per PR advices

* chore: per review

src/common/meta/src/key/table_info.rs

@@ -128,6 +128,7 @@ impl TableInfoValue {
 }
 
 pub type TableInfoManagerRef = Arc<TableInfoManager>;
+#[derive(Clone)]
 pub struct TableInfoManager {
     kv_backend: KvBackendRef,
 }

src/flow/src/adapter.rs

@@ -17,9 +17,16 @@
 
 pub(crate) mod error;
 pub(crate) mod node_context;
+mod table_source;
+mod util;
 
-pub(crate) use node_context::{FlowId, FlownodeContext, TableName};
+pub(crate) use node_context::FlownodeContext;
+pub(crate) use table_source::TableSource;
 
 mod worker;
 
 pub const PER_REQ_MAX_ROW_CNT: usize = 8192;
+// TODO: refactor common types for flow to a separate module
+/// FlowId is a unique identifier for a flow task
+pub type FlowId = u32;
+pub type TableName = [String; 3];

src/flow/src/adapter/error.rs

@@ -98,27 +98,13 @@ pub enum Error {
         location: Location,
     },
 
-    #[snafu(display("Invalid query plan: {source}"))]
-    InvalidQueryPlan {
-        source: query::error::Error,
-        #[snafu(implicit)]
-        location: Location,
-    },
-
     #[snafu(display("Invalid query: prost can't decode substrait plan: {inner}"))]
     InvalidQueryProst {
         inner: api::DecodeError,
         #[snafu(implicit)]
         location: Location,
     },
 
-    #[snafu(display("Invalid query, can't transform to substrait: {source}"))]
-    InvalidQuerySubstrait {
-        source: substrait::error::Error,
-        #[snafu(implicit)]
-        location: Location,
-    },
-
     #[snafu(display("Invalid query: {reason}"))]
     InvalidQuery {
         reason: String,
@@ -193,9 +179,7 @@ impl ErrorExt for Error {
             Self::TableNotFound { .. }
             | Self::TableNotFoundMeta { .. }
             | Self::FlowNotFound { .. } => StatusCode::TableNotFound,
-            Self::InvalidQueryPlan { .. }
-            | Self::InvalidQuerySubstrait { .. }
-            | Self::InvalidQueryProst { .. }
+            Self::InvalidQueryProst { .. }
             | &Self::InvalidQuery { .. }
             | &Self::Plan { .. }
             | &Self::Datatypes { .. } => StatusCode::PlanQuery,

src/flow/src/adapter/node_context.rs

@@ -17,32 +17,18 @@
 use std::collections::{BTreeMap, BTreeSet, HashMap, VecDeque};
 use std::sync::Arc;
 
+use common_telemetry::debug;
 use session::context::QueryContext;
 use snafu::{OptionExt, ResultExt};
 use table::metadata::TableId;
 use tokio::sync::{broadcast, mpsc};
 
 use crate::adapter::error::{Error, EvalSnafu, TableNotFoundSnafu};
+use crate::adapter::{FlowId, TableName, TableSource};
 use crate::expr::error::InternalSnafu;
 use crate::expr::GlobalId;
 use crate::repr::{DiffRow, RelationType, BROADCAST_CAP};
 
-// TODO: refactor common types for flow to a separate module
-/// FlowId is a unique identifier for a flow task
-pub type FlowId = u64;
-pub type TableName = [String; 3];
-
-pub struct TableSource {}
-
-impl TableSource {
-    pub async fn get_table_name_schema(
-        &self,
-        _table_id: &TableId,
-    ) -> Result<(TableName, RelationType), Error> {
-        todo!()
-    }
-}
-
 /// A context that holds the information of the dataflow
 #[derive(Default)]
 pub struct FlownodeContext {
@@ -53,7 +39,7 @@ pub struct FlownodeContext {
     /// broadcast sender for source table, any incoming write request will be sent to the source table's corresponding sender
     ///
     /// Note that we are getting insert requests with table id, so we should use table id as the key
-    pub source_sender: BTreeMap<TableId, broadcast::Sender<DiffRow>>,
+    pub source_sender: BTreeMap<TableId, SourceSender>,
     /// broadcast receiver for sink table, there should only be one receiver, and it will receive all the data from the sink table
     ///
     /// and send it back to the client, since we are mocking the sink table as a client, we should use table name as the key
@@ -74,38 +60,90 @@ pub struct FlownodeContext {
     pub query_context: Option<Arc<QueryContext>>,
 }
 
+/// a simple broadcast sender with backpressure and unbound capacity
+///
+/// receiver still use tokio broadcast channel, since only sender side need to know
+/// backpressure and adjust dataflow running duration to avoid blocking
+pub struct SourceSender {
+    sender: broadcast::Sender<DiffRow>,
+    send_buf: VecDeque<DiffRow>,
+}
+
+impl Default for SourceSender {
+    fn default() -> Self {
+        Self {
+            sender: broadcast::Sender::new(BROADCAST_CAP),
+            send_buf: Default::default(),
+        }
+    }
+}
+
+impl SourceSender {
+    pub fn get_receiver(&self) -> broadcast::Receiver<DiffRow> {
+        self.sender.subscribe()
+    }
+
+    /// send as many as possible rows from send buf
+    /// until send buf is empty or broadchannel is full
+    pub fn try_send_all(&mut self) -> Result<usize, Error> {
+        let mut row_cnt = 0;
+        loop {
+            // if inner sender channel is empty or send buf is empty, there
+            // is nothing to do for now, just break
+            if self.sender.len() >= BROADCAST_CAP || self.send_buf.is_empty() {
+                break;
+            }
+            if let Some(row) = self.send_buf.pop_front() {
+                self.sender
+                    .send(row)
+                    .map_err(|err| {
+                        InternalSnafu {
+                            reason: format!("Failed to send row, error = {:?}", err),
+                        }
+                        .build()
+                    })
+                    .with_context(|_| EvalSnafu)?;
+                row_cnt += 1;
+            }
+        }
+        if row_cnt > 0 {
+            debug!("Send {} rows", row_cnt);
+        }
+
+        Ok(row_cnt)
+    }
+
+    /// return number of rows it actual send(including what's in the buffer)
+    pub fn send_rows(&mut self, rows: Vec<DiffRow>) -> Result<usize, Error> {
+        self.send_buf.extend(rows);
+
+        let row_cnt = self.try_send_all()?;
+
+        Ok(row_cnt)
+    }
+}
+
 impl FlownodeContext {
-    // return number of rows it actual send(including what's in the buffer)
+    /// return number of rows it actual send(including what's in the buffer)
+    ///
+    /// TODO(discord9): make this concurrent
     pub fn send(&mut self, table_id: TableId, rows: Vec<DiffRow>) -> Result<usize, Error> {
         let sender = self
             .source_sender
-            .get(&table_id)
+            .get_mut(&table_id)
             .with_context(|| TableNotFoundSnafu {
                 name: table_id.to_string(),
             })?;
-        let send_buffer = self.send_buffer.entry(table_id).or_default();
-        send_buffer.extend(rows);
-        let mut row_cnt = 0;
-        while let Some(row) = send_buffer.pop_front() {
-            if sender.len() >= BROADCAST_CAP {
-                break;
-            }
-            row_cnt += 1;
-            sender
-                .send(row)
-                .map_err(|err| {
-                    InternalSnafu {
-                        reason: format!(
-                            "Failed to send row to table_id = {:?}, error = {:?}",
-                            table_id, err
-                        ),
-                    }
-                    .build()
-                })
-                .with_context(|_| EvalSnafu)?;
-        }
+        // debug!("FlownodeContext::send: trying to send {} rows", rows.len());
+        sender.send_rows(rows)
+    }
 
-        Ok(row_cnt)
+    /// flush all sender's buf
+    pub fn flush_all_sender(&mut self) -> Result<usize, Error> {
+        self.source_sender
+            .iter_mut()
+            .map(|(_table_id, src_sender)| src_sender.try_send_all())
+            .try_fold(0, |acc, x| x.map(|x| x + acc))
     }
 }
 
@@ -120,7 +158,7 @@ impl FlownodeContext {
         sink_table_name: TableName,
     ) {
         for source_table_id in source_table_ids {
-            self.add_source_sender(*source_table_id);
+            self.add_source_sender_if_not_exist(*source_table_id);
             self.source_to_tasks
                 .entry(*source_table_id)
                 .or_default()
@@ -131,10 +169,9 @@ impl FlownodeContext {
         self.flow_to_sink.insert(task_id, sink_table_name);
     }
 
-    pub fn add_source_sender(&mut self, table_id: TableId) {
-        self.source_sender
-            .entry(table_id)
-            .or_insert_with(|| broadcast::channel(BROADCAST_CAP).0);
+    /// try add source sender, if already exist, do nothing
+    pub fn add_source_sender_if_not_exist(&mut self, table_id: TableId) {
+        let _sender = self.source_sender.entry(table_id).or_default();
     }
 
     pub fn add_sink_receiver(&mut self, table_name: TableName) {
@@ -143,10 +180,7 @@ impl FlownodeContext {
             .or_insert_with(mpsc::unbounded_channel::<DiffRow>);
     }
 
-    pub fn get_source_by_global_id(
-        &self,
-        id: &GlobalId,
-    ) -> Result<&broadcast::Sender<DiffRow>, Error> {
+    pub fn get_source_by_global_id(&self, id: &GlobalId) -> Result<&SourceSender, Error> {
         let table_id = self
             .table_repr
             .get_by_global_id(id)