coprocessor: collect output counts for each executor

src/coprocessor/dag/dag.rs

@@ -34,6 +34,7 @@ pub struct DAGContext {
     exec: Box<Executor>,
     output_offsets: Vec<u32>,
     batch_row_limit: usize,
+    num_execs: usize,
 }
 
 impl DAGContext {
@@ -63,6 +64,7 @@ impl DAGContext {
             exec: dag_executor.exec,
             output_offsets: req.take_output_offsets(),
             batch_row_limit: batch_row_limit,
+            num_execs: req.take_executors().len(),
         })
     }
 
@@ -91,6 +93,9 @@ impl DAGContext {
                     let mut resp = Response::new();
                     let mut sel_resp = SelectResponse::new();
                     sel_resp.set_chunks(RepeatedField::from_vec(chunks));
+                    let mut counts = Vec::with_capacity(self.num_execs);
+                    self.exec.collect_output_counts(&mut counts);
+                    sel_resp.set_output_counts(counts);
                     let data = box_try!(sel_resp.write_to_bytes());
                     resp.set_data(data);
                     return Ok(resp);

src/coprocessor/dag/executor/aggregation.rs

@@ -80,6 +80,7 @@ pub struct AggregationExecutor {
     cols: Arc<Vec<ColumnInfo>>,
     related_cols_offset: Vec<usize>, // offset of related columns
     src: Box<Executor>,
+    count: i64,
 }
 
 impl AggregationExecutor {
@@ -108,6 +109,7 @@ impl AggregationExecutor {
             cols: columns,
             related_cols_offset: visitor.column_offsets(),
             src: src,
+            count: 0,
         })
     }
 
@@ -159,6 +161,7 @@ impl AggregationExecutor {
 
 impl Executor for AggregationExecutor {
     fn next(&mut self) -> Result<Option<Row>> {
+        self.count += 1;
         if !self.executed {
             self.aggregate()?;
             self.executed = true;
@@ -190,6 +193,11 @@ impl Executor for AggregationExecutor {
         }
     }
 
+    fn collect_output_counts(&mut self, counts: &mut Vec<i64>) {
+        self.src.collect_output_counts(counts);
+        counts.push(self.count - 1);
+    }
+
     fn collect_statistics_into(&mut self, statistics: &mut Statistics) {
         self.src.collect_statistics_into(statistics);
     }
@@ -364,5 +372,9 @@ mod test {
             assert_eq!(ds[3], Datum::from(expect_cols.3));
             assert_eq!(ds[4], Datum::from(expect_cols.4));
         }
+        let expected_counts = vec![raw_data.len() as i64, expect_row_cnt as i64];
+        let mut counts = Vec::with_capacity(2);
+        aggr_ect.collect_output_counts(&mut counts);
+        assert_eq!(expected_counts, counts);
     }
 }

src/coprocessor/dag/executor/index_scan.rs

@@ -37,6 +37,7 @@ pub struct IndexScanExecutor {
     key_ranges: IntoIter<KeyRange>,
     scanner: Option<Scanner>,
     unique: bool,
+    count: i64,
 }
 
 impl IndexScanExecutor {
@@ -67,6 +68,7 @@ impl IndexScanExecutor {
             key_ranges: key_ranges.into_iter(),
             scanner: None,
             unique: unique,
+            count: 0,
         }
     }
 
@@ -86,6 +88,7 @@ impl IndexScanExecutor {
             key_ranges: key_ranges.into_iter(),
             scanner: None,
             unique: false,
+            count: 0,
         }
     }
 
@@ -145,6 +148,7 @@ impl IndexScanExecutor {
 
 impl Executor for IndexScanExecutor {
     fn next(&mut self) -> Result<Option<Row>> {
+        self.count += 1;
         loop {
             if let Some(row) = self.get_row_from_range_scanner()? {
                 return Ok(Some(row));
@@ -170,6 +174,10 @@ impl Executor for IndexScanExecutor {
         }
     }
 
+    fn collect_output_counts(&mut self, counts: &mut Vec<i64>) {
+        counts.push(self.count - 1);
+    }
+
     fn collect_statistics_into(&mut self, statistics: &mut Statistics) {
         statistics.add(&self.statistics);
         self.statistics = Statistics::default();
@@ -360,6 +368,10 @@ mod test {
             }
         }
         assert!(scanner.next().unwrap().is_none());
+        let expected_counts = vec![KEY_NUMBER as i64];
+        let mut counts = Vec::with_capacity(1);
+        scanner.collect_output_counts(&mut counts);
+        assert_eq!(expected_counts, counts);
     }
 
     #[test]

src/coprocessor/dag/executor/limit.rs

@@ -51,6 +51,10 @@ impl<'a> Executor for LimitExecutor<'a> {
         }
     }
 
+    fn collect_output_counts(&mut self, _: &mut Vec<i64>) {
+        // We do not know whether `limit` has consumed all of it's source, so just ignore it.
+    }
+
     fn collect_statistics_into(&mut self, statistics: &mut Statistics) {
         self.src.collect_statistics_into(statistics);
     }

src/coprocessor/dag/executor/mod.rs

@@ -129,6 +129,7 @@ impl Row {
 
 pub trait Executor {
     fn next(&mut self) -> Result<Option<Row>>;
+    fn collect_output_counts(&mut self, counts: &mut Vec<i64>);
     fn collect_statistics_into(&mut self, stats: &mut Statistics);
 }
 

src/coprocessor/dag/executor/selection.rs

@@ -30,6 +30,7 @@ pub struct SelectionExecutor {
     related_cols_offset: Vec<usize>, // offset of related columns
     ctx: Arc<EvalContext>,
     src: Box<Executor>,
+    count: i64,
 }
 
 impl SelectionExecutor {
@@ -49,13 +50,15 @@ impl SelectionExecutor {
             related_cols_offset: visitor.column_offsets(),
             ctx: ctx,
             src: src,
+            count: 0,
         })
     }
 }
 
 #[allow(never_loop)]
 impl Executor for SelectionExecutor {
     fn next(&mut self) -> Result<Option<Row>> {
+        self.count += 1;
         'next: while let Some(row) = self.src.next()? {
             let cols = inflate_with_col_for_dag(
                 &self.ctx,
@@ -75,6 +78,11 @@ impl Executor for SelectionExecutor {
         Ok(None)
     }
 
+    fn collect_output_counts(&mut self, counts: &mut Vec<i64>) {
+        self.src.collect_output_counts(counts);
+        counts.push(self.count - 1);
+    }
+
     fn collect_statistics_into(&mut self, statistics: &mut Statistics) {
         self.src.collect_statistics_into(statistics);
     }
@@ -276,5 +284,9 @@ mod tests {
         assert_eq!(selection_rows.len(), expect_row_handles.len());
         let result_row = selection_rows.iter().map(|r| r.handle).collect::<Vec<_>>();
         assert_eq!(result_row, expect_row_handles);
+        let expected_counts = vec![raw_data.len() as i64, expect_row_handles.len() as i64];
+        let mut counts = Vec::with_capacity(2);
+        selection_executor.collect_output_counts(&mut counts);
+        assert_eq!(expected_counts, counts);
     }
 }

src/coprocessor/dag/executor/table_scan.rs

@@ -34,6 +34,7 @@ pub struct TableScanExecutor {
     col_ids: HashSet<i64>,
     key_ranges: IntoIter<KeyRange>,
     scanner: Option<Scanner>,
+    count: i64,
 }
 
 impl TableScanExecutor {
@@ -62,6 +63,7 @@ impl TableScanExecutor {
             col_ids: col_ids,
             key_ranges: key_ranges.into_iter(),
             scanner: None,
+            count: 0,
         }
     }
 
@@ -97,6 +99,7 @@ impl TableScanExecutor {
 
 impl Executor for TableScanExecutor {
     fn next(&mut self) -> Result<Option<Row>> {
+        self.count += 1;
         loop {
             if let Some(row) = self.get_row_from_range_scanner()? {
                 return Ok(Some(row));
@@ -123,6 +126,10 @@ impl Executor for TableScanExecutor {
         }
     }
 
+    fn collect_output_counts(&mut self, counts: &mut Vec<i64>) {
+        counts.push(self.count - 1);
+    }
+
     fn collect_statistics_into(&mut self, statistics: &mut Statistics) {
         statistics.add(&self.statistics);
         self.statistics = Statistics::default();
@@ -210,6 +217,10 @@ mod test {
             assert_eq!(expect_row[&cid], v.to_vec());
         }
         assert!(table_scanner.next().unwrap().is_none());
+        let expected_counts = vec![1];
+        let mut counts = Vec::with_capacity(1);
+        table_scanner.collect_output_counts(&mut counts);
+        assert_eq!(expected_counts, counts);
     }
 
     #[test]

src/coprocessor/dag/executor/topn.rs

@@ -62,6 +62,7 @@ pub struct TopNExecutor {
     iter: Option<IntoIter<SortRow>>,
     ctx: Arc<EvalContext>,
     src: Box<Executor>,
+    count: i64,
 }
 
 impl TopNExecutor {
@@ -87,6 +88,7 @@ impl TopNExecutor {
             iter: None,
             ctx: ctx,
             src: src,
+            count: 0,
         })
     }
 
@@ -114,6 +116,7 @@ impl TopNExecutor {
 
 impl Executor for TopNExecutor {
     fn next(&mut self) -> Result<Option<Row>> {
+        self.count += 1;
         if self.iter.is_none() {
             self.fetch_all()?;
             self.iter = Some(self.heap.take().unwrap().into_sorted_vec()?.into_iter());
@@ -128,6 +131,11 @@ impl Executor for TopNExecutor {
         }
     }
 
+    fn collect_output_counts(&mut self, counts: &mut Vec<i64>) {
+        self.src.collect_output_counts(counts);
+        counts.push(self.count - 1);
+    }
+
     fn collect_statistics_into(&mut self, statistics: &mut Statistics) {
         self.src.collect_statistics_into(statistics);
     }
@@ -436,5 +444,9 @@ pub mod test {
         for (row, handle) in topn_rows.iter().zip(expect_row_handles) {
             assert_eq!(row.handle, handle);
         }
+        let expected_counts = vec![6, limit as i64];
+        let mut counts = Vec::with_capacity(2);
+        topn_ect.collect_output_counts(&mut counts);
+        assert_eq!(expected_counts, counts);
     }
 }

tests/coprocessor/test_select.rs

@@ -2730,3 +2730,22 @@ fn test_key_is_locked_for_index() {
     assert!(resp.has_locked(), "{:?}", resp);
     end_point.stop().unwrap().join().unwrap();
 }
+
+#[test]
+fn test_output_counts() {
+    let data = vec![
+        (1, Some("name:0"), 2),
+        (2, Some("name:4"), 3),
+        (4, Some("name:3"), 1),
+        (5, Some("name:1"), 4),
+    ];
+
+    let product = ProductTable::new();
+    let (_, mut end_point) = init_with_data(&product, &data);
+
+    let req = DAGSelect::from(&product.table).build();
+    let resp = handle_select(&end_point, req);
+    assert_eq!(resp.get_output_counts(), [data.len() as i64]);
+
+    end_point.stop().unwrap().join().unwrap();
+}