add indexed variable in PhysicalScalar

Author: leiysky

src/query/service/src/api/rpc/flight_scatter_hash_v2.rs

@@ -27,13 +27,12 @@ use crate::api::rpc::flight_scatter::FlightScatter;
 use crate::evaluator::EvalNode;
 use crate::evaluator::Evaluator;
 use crate::evaluator::TypedVector;
-use crate::sql::executor::ColumnID;
 use crate::sql::executor::PhysicalScalar;
 
 #[derive(Clone)]
 pub struct HashFlightScatterV2 {
     func_ctx: FunctionContext,
-    hash_keys: Vec<EvalNode<ColumnID>>,
+    hash_keys: Vec<EvalNode>,
     hash_functions: Vec<Box<dyn Function>>,
     scatter_size: usize,
 }
@@ -118,7 +117,7 @@ impl HashFlightScatterV2 {
 struct OneHashKeyFlightScatter {
     scatter_size: usize,
     func_ctx: FunctionContext,
-    indices_scalar: EvalNode<ColumnID>,
+    indices_scalar: EvalNode,
 }
 
 impl OneHashKeyFlightScatter {
@@ -129,7 +128,7 @@ impl OneHashKeyFlightScatter {
     ) -> Result<Box<dyn FlightScatter>> {
         let hash_key = Evaluator::eval_physical_scalar(scalar)?;
 
-        let mut sip_hash = EvalNode::<ColumnID>::Function {
+        let mut sip_hash = EvalNode::Function {
             args: vec![hash_key],
             func: FunctionFactory::instance().get("sipHash", &[&scalar.data_type()])?,
         };
@@ -146,7 +145,7 @@ impl OneHashKeyFlightScatter {
         Ok(Box::new(OneHashKeyFlightScatter {
             scatter_size,
             func_ctx,
-            indices_scalar: EvalNode::<ColumnID>::Function {
+            indices_scalar: EvalNode::Function {
                 args: vec![sip_hash, EvalNode::Constant {
                     value: DataValue::UInt64(scatter_size as u64),
                     data_type: DataTypeImpl::UInt64(UInt64Type::new()),

src/query/service/src/evaluator/eval_context.rs

@@ -12,8 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::fmt::Debug;
-
+use common_datablocks::DataBlock;
 use common_datavalues::ColumnWithField;
 use common_datavalues::DataField;
 use common_datavalues::DataTypeImpl;
@@ -22,14 +21,12 @@ use common_exception::Result;
 use common_functions::scalars::Function;
 use common_functions::scalars::FunctionContext;
 
-use crate::evaluator::eval_context::EmptyEvalContext;
-use crate::evaluator::EvalContext;
 use crate::evaluator::TypedVector;
 
 /// A intermediate representation of a evaluable scalar expression, with configurable
 /// EvalContext.
 #[derive(Clone)]
-pub enum EvalNode<VectorID> {
+pub enum EvalNode {
     Function {
         func: Box<dyn Function>,
         args: Vec<Self>,
@@ -41,48 +38,63 @@ pub enum EvalNode<VectorID> {
         data_type: DataTypeImpl,
     },
     Variable {
-        id: VectorID,
+        name: String,
+    },
+    IndexedVariable {
+        index: usize,
     },
 }
 
-impl<VectorID> EvalNode<VectorID>
-where VectorID: PartialEq + Eq + Clone + Debug
-{
-    /// Evaluate with given context, which is typically a `DataBlock`
-    pub fn eval(
-        &self,
-        func_ctx: &FunctionContext,
-        eval_ctx: &impl EvalContext<VectorID = VectorID>,
-    ) -> Result<TypedVector> {
+impl EvalNode {
+    pub fn eval(&self, func_ctx: &FunctionContext, data_block: &DataBlock) -> Result<TypedVector> {
         match &self {
             EvalNode::Function { func, args } => {
                 let args = args
                     .iter()
                     .map(|arg| {
-                        let vector = arg.eval(func_ctx, eval_ctx)?;
+                        let vector = arg.eval(func_ctx, data_block)?;
                         Ok(ColumnWithField::new(
                             vector.vector,
                             DataField::new("", vector.logical_type),
                         ))
                     })
                     .collect::<Result<Vec<_>>>()?;
                 Ok(TypedVector::new(
-                    func.eval(func_ctx.clone(), &args, eval_ctx.tuple_count())?,
+                    func.eval(func_ctx.clone(), &args, data_block.num_rows())?,
                     func.return_type(),
                 ))
             }
             EvalNode::Constant { value, data_type } => {
-                let vector = value.as_const_column(data_type, eval_ctx.tuple_count())?;
+                let vector = value.as_const_column(data_type, data_block.num_rows())?;
                 Ok(TypedVector::new(vector, data_type.clone()))
             }
-            EvalNode::Variable { id } => eval_ctx.get_vector(id),
+            EvalNode::Variable { name } => {
+                let column = data_block.try_column_by_name(name)?;
+                let data_type = data_block
+                    .schema()
+                    .field_with_name(name)?
+                    .data_type()
+                    .clone();
+                Ok(TypedVector {
+                    vector: column.clone(),
+                    logical_type: data_type,
+                })
+            }
+            EvalNode::IndexedVariable { index } => {
+                let column = data_block.column(*index);
+                let data_type = data_block.schema().field(*index).data_type().clone();
+                Ok(TypedVector {
+                    vector: column.clone(),
+                    logical_type: data_type,
+                })
+            }
         }
     }
 
     /// Try to evaluate as a constant expression
     pub fn try_eval_const(&self, func_ctx: &FunctionContext) -> Result<(DataValue, DataTypeImpl)> {
-        let eval_ctx = EmptyEvalContext::<VectorID>::new();
-        let vector = self.eval(func_ctx, &eval_ctx)?;
+        let dummy_data_block = DataBlock::empty();
+        let vector = self.eval(func_ctx, &dummy_data_block)?;
         debug_assert!(vector.vector.len() == 1);
         Ok((vector.vector.get(0), vector.logical_type))
     }

src/query/service/src/evaluator/eval_node.rs

@@ -12,13 +12,39 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-mod eval_context;
 mod eval_node;
 mod physical_scalar;
 mod scalar;
 
-pub use eval_context::EvalContext;
-pub use eval_context::TypedVector;
+use common_datavalues::ColumnRef;
+use common_datavalues::DataTypeImpl;
 pub use eval_node::EvalNode;
 
 pub struct Evaluator;
+
+#[derive(Clone, Debug)]
+pub struct TypedVector {
+    pub(super) vector: ColumnRef,
+    pub(super) logical_type: DataTypeImpl,
+}
+
+impl TypedVector {
+    pub fn new(data: ColumnRef, logical_type: DataTypeImpl) -> Self {
+        Self {
+            vector: data,
+            logical_type,
+        }
+    }
+
+    pub fn logical_type(&self) -> DataTypeImpl {
+        self.logical_type.clone()
+    }
+
+    pub fn physical_type(&self) -> DataTypeImpl {
+        self.vector.data_type()
+    }
+
+    pub fn vector(&self) -> &ColumnRef {
+        &self.vector
+    }
+}

src/query/service/src/evaluator/mod.rs

@@ -23,23 +23,17 @@ use crate::evaluator::Evaluator;
 use crate::sql::executor::PhysicalScalar;
 
 impl Evaluator {
-    pub fn eval_physical_scalars<VectorID>(
-        physical_scalars: &[PhysicalScalar],
-    ) -> Result<Vec<EvalNode<VectorID>>>
-    where VectorID: From<String> {
+    pub fn eval_physical_scalars(physical_scalars: &[PhysicalScalar]) -> Result<Vec<EvalNode>> {
         physical_scalars
             .iter()
-            .map(Evaluator::eval_physical_scalar::<VectorID>)
+            .map(Evaluator::eval_physical_scalar)
             .collect::<Result<_>>()
     }
 
-    pub fn eval_physical_scalar<VectorID>(
-        physical_scalar: &PhysicalScalar,
-    ) -> Result<EvalNode<VectorID>>
-    where VectorID: From<String> {
+    pub fn eval_physical_scalar(physical_scalar: &PhysicalScalar) -> Result<EvalNode> {
         match physical_scalar {
             PhysicalScalar::Variable { column_id, .. } => Ok(EvalNode::Variable {
-                id: column_id.clone().into(),
+                name: column_id.clone(),
             }),
             PhysicalScalar::Constant { value, data_type } => Ok(EvalNode::Constant {
                 value: value.clone(),
@@ -66,6 +60,9 @@ impl Evaluator {
                     args: vec![Self::eval_physical_scalar(input)?],
                 })
             }
+            PhysicalScalar::IndexedVariable { index, .. } => {
+                Ok(EvalNode::IndexedVariable { index: *index })
+            }
         }
     }
 }

src/query/service/src/evaluator/physical_scalar.rs

@@ -25,11 +25,10 @@ use crate::sql::plans::Scalar;
 use crate::sql::plans::ScalarExpr;
 
 impl Evaluator {
-    pub fn eval_scalar<VectorID>(scalar: &Scalar) -> Result<EvalNode<VectorID>>
-    where VectorID: From<String> {
+    pub fn eval_scalar(scalar: &Scalar) -> Result<EvalNode> {
         match scalar {
             Scalar::BoundColumnRef(column_ref) => Ok(EvalNode::Variable {
-                id: column_ref.column.index.to_string().into(),
+                name: column_ref.column.index.to_string(),
             }),
             Scalar::ConstantExpr(constant) => Ok(EvalNode::Constant {
                 value: constant.value.clone(),
@@ -62,7 +61,7 @@ impl Evaluator {
                 Ok(EvalNode::Function { func, args })
             }
             Scalar::FunctionCall(func) => {
-                let args: Vec<EvalNode<VectorID>> = func
+                let args: Vec<EvalNode> = func
                     .arguments
                     .iter()
                     .map(Self::eval_scalar)

src/query/service/src/evaluator/scalar.rs

@@ -19,16 +19,15 @@ use parking_lot::RwLock;
 use crate::evaluator::EvalNode;
 use crate::evaluator::Evaluator;
 use crate::pipelines::processors::transforms::hash_join::MarkJoinDesc;
-use crate::sql::executor::ColumnID;
 use crate::sql::executor::HashJoin;
 use crate::sql::executor::PhysicalScalar;
 use crate::sql::plans::JoinType;
 
 pub struct HashJoinDesc {
-    pub(crate) build_keys: Vec<EvalNode<ColumnID>>,
-    pub(crate) probe_keys: Vec<EvalNode<ColumnID>>,
+    pub(crate) build_keys: Vec<EvalNode>,
+    pub(crate) probe_keys: Vec<EvalNode>,
     pub(crate) join_type: JoinType,
-    pub(crate) other_predicate: Option<EvalNode<ColumnID>>,
+    pub(crate) other_predicate: Option<EvalNode>,
     pub(crate) marker_join_desc: MarkJoinDesc,
     /// Whether the Join are derived from correlated subquery.
     pub(crate) from_correlated_subquery: bool,

src/query/service/src/pipelines/processors/transforms/hash_join/desc.rs

@@ -41,7 +41,6 @@ use crate::evaluator::EvalNode;
 use crate::pipelines::processors::transforms::hash_join::join_hash_table::MarkerKind;
 use crate::pipelines::processors::transforms::hash_join::row::RowPtr;
 use crate::sessions::TableContext;
-use crate::sql::executor::ColumnID;
 use crate::sql::planner::plans::JoinType;
 use crate::sql::plans::JoinType::Mark;
 
@@ -602,7 +601,7 @@ impl JoinHashTable {
     fn get_other_filters(
         &self,
         merged_block: &DataBlock,
-        filter: &EvalNode<ColumnID>,
+        filter: &EvalNode,
     ) -> Result<(Option<Bitmap>, bool, bool)> {
         let func_ctx = self.ctx.try_get_function_context()?;
         // `predicate_column` contains a column, which is a boolean column.