coprocessor: Batch Stream Aggregation Executor

src/coprocessor/codec/data_type/mod.rs

@@ -14,7 +14,7 @@ pub type Bytes = Vec<u8>;
 pub use crate::coprocessor::codec::mysql::{Decimal, Duration, Json, Time as DateTime};
 
 // Dynamic eval types.
-pub use self::scalar::ScalarValue;
+pub use self::scalar::{ScalarValue, ScalarValueRef};
 pub use self::vector::{VectorValue, VectorValueExt};
 pub use self::vector_like::{VectorLikeValueRef, VectorLikeValueRefSpecialized};
 

src/coprocessor/codec/data_type/scalar.rs

@@ -28,6 +28,40 @@ pub enum ScalarValue {
     Json(Option<super::Json>),
 }
 
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum ScalarValueRef<'a> {
+    Int(&'a Option<super::Int>),
+    Real(&'a Option<super::Real>),
+    Decimal(&'a Option<super::Decimal>),
+    Bytes(&'a Option<super::Bytes>),
+    DateTime(&'a Option<super::DateTime>),
+    Duration(&'a Option<super::Duration>),
+    Json(&'a Option<super::Json>),
+}
+
+impl<'a> ScalarValueRef<'a> {
+    #[inline]
+    #[allow(clippy::clone_on_copy)]
+    pub fn to_owned(self) -> ScalarValue {
+        match_template_evaluable! {
+            TT, match self {
+                ScalarValueRef::TT(v) => ScalarValue::TT(v.clone()),
+            }
+        }
+    }
+}
+
+impl<'a> PartialEq<ScalarValue> for ScalarValueRef<'a> {
+    fn eq(&self, other: &ScalarValue) -> bool {
+        match_template_evaluable! {
+            TT, match (self, other) {
+                (ScalarValueRef::TT(v1), ScalarValue::TT(v2)) => v1 == &v2,
+                _ => false
+            }
+        }
+    }
+}
+
 impl ScalarValue {
     #[inline]
     pub fn eval_type(&self) -> EvalType {

src/coprocessor/codec/data_type/vector.rs

@@ -6,6 +6,7 @@ use cop_datatype::{EvalType, FieldTypeAccessor, FieldTypeFlag, FieldTypeTp};
 use tipb::expression::FieldType;
 
 use super::*;
+use crate::coprocessor::codec::data_type::scalar::ScalarValueRef;
 use crate::coprocessor::codec::datum;
 use crate::coprocessor::codec::mysql::Tz;
 use crate::coprocessor::codec::{Error, Result};
@@ -175,6 +176,20 @@ impl VectorValue {
         Ok(())
     }
 
+    /// Returns a `ScalarValueRef` to the element at the index.
+    ///
+    /// # Panics
+    ///
+    /// Panics if index is out of range.
+    #[inline]
+    pub fn get_scalar_ref(&self, index: usize) -> ScalarValueRef<'_> {
+        match_template_evaluable! {
+            TT, match self {
+                VectorValue::TT(v) => ScalarValueRef::TT(&v[index]),
+            }
+        }
+    }
+
     /// Pushes a value into the column by decoding the datum and converting to current
     /// column's type.
     ///

src/coprocessor/dag/batch/executors/fast_hash_aggr_executor.rs

@@ -262,11 +262,14 @@ impl<Src: BatchExecutor> AggregationExecutorImpl<Src> for FastHashAggregationImp
     }
 
     #[inline]
-    fn iterate_each_group_for_aggregation(
+    fn iterate_available_groups(
         &mut self,
         entities: &mut Entities<Src>,
+        src_is_drained: bool,
         mut iteratee: impl FnMut(&mut Entities<Src>, &[Box<dyn AggrFunctionState>]) -> Result<()>,
     ) -> Result<Vec<LazyBatchColumn>> {
+        assert!(src_is_drained);
+
         let aggr_fns_len = entities.each_aggr_fn.len();
         let mut group_by_column = LazyBatchColumn::decoded_with_capacity_and_tp(
             self.groups.len(),
@@ -290,6 +293,12 @@ impl<Src: BatchExecutor> AggregationExecutorImpl<Src> for FastHashAggregationImp
 
         Ok(vec![group_by_column])
     }
+
+    /// Fast hash aggregation can output aggregate results only if the source is drained.
+    #[inline]
+    fn is_partial_results_ready(&self) -> bool {
+        false
+    }
 }
 
 fn calc_groups_each_row<T: Evaluable + Eq + std::hash::Hash>(

src/coprocessor/dag/batch/executors/mod.rs

@@ -6,6 +6,7 @@ mod limit_executor;
 mod selection_executor;
 mod simple_aggr_executor;
 mod slow_hash_aggr_executor;
+mod stream_aggr_executor;
 mod table_scan_executor;
 mod util;
 
@@ -15,4 +16,5 @@ pub use self::limit_executor::BatchLimitExecutor;
 pub use self::selection_executor::BatchSelectionExecutor;
 pub use self::simple_aggr_executor::BatchSimpleAggregationExecutor;
 pub use self::slow_hash_aggr_executor::BatchSlowHashAggregationExecutor;
+pub use self::stream_aggr_executor::BatchStreamAggregationExecutor;
 pub use self::table_scan_executor::BatchTableScanExecutor;

src/coprocessor/dag/batch/executors/simple_aggr_executor.rs

@@ -163,14 +163,22 @@ impl<Src: BatchExecutor> AggregationExecutorImpl<Src> for SimpleAggregationImpl
     }
 
     #[inline]
-    fn iterate_each_group_for_aggregation(
+    fn iterate_available_groups(
         &mut self,
         entities: &mut Entities<Src>,
+        src_is_drained: bool,
         mut iteratee: impl FnMut(&mut Entities<Src>, &[Box<dyn AggrFunctionState>]) -> Result<()>,
     ) -> Result<Vec<LazyBatchColumn>> {
+        assert!(src_is_drained);
         iteratee(entities, &self.states)?;
         Ok(Vec::new())
     }
+
+    /// Simple aggregation can output aggregate results only if the source is drained.
+    #[inline]
+    fn is_partial_results_ready(&self) -> bool {
+        false
+    }
 }
 
 #[cfg(test)]

src/coprocessor/dag/batch/executors/slow_hash_aggr_executor.rs

@@ -255,11 +255,14 @@ impl<Src: BatchExecutor> AggregationExecutorImpl<Src> for SlowHashAggregationImp
     }
 
     #[inline]
-    fn iterate_each_group_for_aggregation(
+    fn iterate_available_groups(
         &mut self,
         entities: &mut Entities<Src>,
+        src_is_drained: bool,
         mut iteratee: impl FnMut(&mut Entities<Src>, &[Box<dyn AggrFunctionState>]) -> Result<()>,
     ) -> Result<Vec<LazyBatchColumn>> {
+        assert!(src_is_drained);
+
         let number_of_groups = self.groups.len();
         let group_by_exps_len = self.group_by_exps.len();
         let mut group_by_columns: Vec<_> = self
@@ -287,6 +290,12 @@ impl<Src: BatchExecutor> AggregationExecutorImpl<Src> for SlowHashAggregationImp
 
         Ok(group_by_columns)
     }
+
+    /// Slow hash aggregation can output aggregate results only if the source is drained.
+    #[inline]
+    fn is_partial_results_ready(&self) -> bool {
+        false
+    }
 }
 
 #[cfg(test)]