chore(query): refactoring ValueType, pass a data type parameter to upcast method (#18082)

Author: forsaken628

src/common/native/src/read/array/double.rs

@@ -17,9 +17,9 @@ use std::io::Cursor;
 use std::marker::PhantomData;
 
 use databend_common_column::buffer::Buffer;
+use databend_common_expression::types::ArgType;
 use databend_common_expression::types::Number;
 use databend_common_expression::types::NumberType;
-use databend_common_expression::types::ValueType;
 use databend_common_expression::Column;
 use databend_common_expression::TableDataType;
 

src/common/native/src/read/array/interval.rs

@@ -16,8 +16,8 @@ use std::io::Cursor;
 
 use databend_common_column::buffer::Buffer;
 use databend_common_column::types::months_days_micros;
+use databend_common_expression::types::ArgType;
 use databend_common_expression::types::IntervalType;
-use databend_common_expression::types::ValueType;
 use databend_common_expression::Column;
 use databend_common_expression::TableDataType;
 

src/common/native/src/write/serialize.rs

@@ -18,6 +18,7 @@ use databend_common_column::bitmap::Bitmap;
 use databend_common_column::buffer::Buffer;
 use databend_common_column::types::i256;
 use databend_common_expression::types::AnyType;
+use databend_common_expression::types::DataType;
 use databend_common_expression::types::DecimalColumn;
 use databend_common_expression::types::GeographyColumn;
 use databend_common_expression::types::NullableColumn;
@@ -78,7 +79,7 @@ impl<'a, W: Write> ValueVisitor for WriteVisitor<'a, W> {
         unreachable!()
     }
 
-    fn visit_typed_column<T: ValueType>(&mut self, _column: T::Column) -> Result<()> {
+    fn visit_typed_column<T: ValueType>(&mut self, _: T::Column, _: &DataType) -> Result<()> {
         unreachable!()
     }
 

src/query/catalog/src/plan/internal_column.rs

@@ -15,7 +15,6 @@
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_expression::types::number::F32;
-use databend_common_expression::types::string::StringColumnBuilder;
 use databend_common_expression::types::DataType;
 use databend_common_expression::types::DecimalDataType;
 use databend_common_expression::types::DecimalSize;
@@ -226,35 +225,17 @@ impl InternalColumn {
                     }
                 }
 
-                BlockEntry::new(
-                    DataType::Number(NumberDataType::UInt64),
-                    Value::Column(UInt64Type::from_data(row_ids)),
-                )
+                UInt64Type::from_data(row_ids).into()
             }
             InternalColumnType::BlockName => {
-                let mut builder = StringColumnBuilder::with_capacity(1);
-                builder.put_and_commit(&meta.block_location);
-                BlockEntry::new(
-                    DataType::String,
-                    Value::Scalar(Scalar::String(builder.build_scalar())),
-                )
+                BlockEntry::from_arg_scalar::<StringType>(meta.block_location.clone())
             }
             InternalColumnType::SegmentName => {
-                let mut builder = StringColumnBuilder::with_capacity(1);
-                builder.put_and_commit(&meta.segment_location);
-                BlockEntry::new(
-                    DataType::String,
-                    Value::Scalar(Scalar::String(builder.build_scalar())),
-                )
-            }
-            InternalColumnType::SnapshotName => {
-                let mut builder = StringColumnBuilder::with_capacity(1);
-                builder.put_and_commit(meta.snapshot_location.clone().unwrap_or("".to_string()));
-                BlockEntry::new(
-                    DataType::String,
-                    Value::Scalar(Scalar::String(builder.build_scalar())),
-                )
+                BlockEntry::from_arg_scalar::<StringType>(meta.segment_location.clone())
             }
+            InternalColumnType::SnapshotName => BlockEntry::from_arg_scalar::<StringType>(
+                meta.snapshot_location.clone().unwrap_or_default(),
+            ),
             InternalColumnType::BaseRowId => {
                 let uuid =
                     try_extract_uuid_str_from_path(&meta.block_location).unwrap_or_else(|e| {
@@ -275,10 +256,7 @@ impl InternalColumn {
                         row_ids.push(row_id);
                     }
                 }
-                BlockEntry::new(
-                    DataType::String,
-                    Value::Column(StringType::from_data(row_ids)),
-                )
+                StringType::from_data(row_ids).into()
             }
             InternalColumnType::BaseBlockIds => {
                 assert!(meta.base_block_ids.is_some());
@@ -295,10 +273,7 @@ impl InternalColumn {
                 for (idx, _) in matched_rows.iter() {
                     bitmap.set(*idx, true);
                 }
-                BlockEntry::new(
-                    DataType::Boolean,
-                    Value::Column(Column::Boolean(bitmap.into())),
-                )
+                Column::Boolean(bitmap.into()).into()
             }
             InternalColumnType::SearchScore => {
                 assert!(meta.matched_rows.is_some());
@@ -311,10 +286,7 @@ impl InternalColumn {
                         *val = F32::from(*score.unwrap());
                     }
                 }
-                BlockEntry::new(
-                    DataType::Number(NumberDataType::Float32),
-                    Value::Column(Float32Type::from_data(scores)),
-                )
+                Float32Type::from_data(scores).into()
             }
             InternalColumnType::FileName | InternalColumnType::FileRowNumber => {
                 todo!("generate_column_values not support for file related")

src/query/expression/src/aggregate/group_hash.rs

@@ -27,6 +27,7 @@ use crate::types::BinaryColumn;
 use crate::types::BinaryType;
 use crate::types::BitmapType;
 use crate::types::BooleanType;
+use crate::types::DataType;
 use crate::types::DateType;
 use crate::types::Decimal64As128Type;
 use crate::types::DecimalColumn;
@@ -85,8 +86,14 @@ impl<const IS_FIRST: bool> ValueVisitor for HashVisitor<'_, IS_FIRST> {
         unreachable!()
     }
 
-    fn visit_typed_column<T: ValueType>(&mut self, column: T::Column) -> Result<()> {
-        self.combine_group_hash_type_column::<AnyType>(&T::upcast_column(column));
+    fn visit_typed_column<T: ValueType>(
+        &mut self,
+        column: T::Column,
+        data_type: &DataType,
+    ) -> Result<()> {
+        self.combine_group_hash_type_column::<AnyType>(&T::upcast_column_with_type(
+            column, data_type,
+        ));
         Ok(())
     }
 
@@ -386,11 +393,16 @@ where I: Index
         }
     }
 
-    fn visit_typed_column<T: ValueType>(&mut self, column: T::Column) -> Result<()> {
+    fn visit_typed_column<T: ValueType>(
+        &mut self,
+        column: T::Column,
+        data_type: &DataType,
+    ) -> Result<()> {
         self.visit_indices(|i| {
-            let x = T::upcast_scalar(T::to_owned_scalar(
-                T::index_column(&column, i.to_usize()).unwrap(),
-            ));
+            let x = T::upcast_scalar_with_type(
+                T::to_owned_scalar(T::index_column(&column, i.to_usize()).unwrap()),
+                data_type,
+            );
             x.as_ref().agg_hash()
         })
     }
@@ -628,8 +640,11 @@ mod tests {
 
         {
             let c = Int32Type::from_data(vec![3, 1, 2]);
-            let c = NullableColumn::<AnyType>::new(c, Bitmap::from([true, true, false]));
-            let nc = NullableType::<AnyType>::upcast_column(c);
+            let c = NullableColumn::new(c, Bitmap::from([true, true, false]));
+            let nc = NullableType::<AnyType>::upcast_column_with_type(
+                c,
+                &Int32Type::data_type().wrap_nullable(),
+            );
 
             let indices = [0, 1, 2];
             let mut target = vec![0; 3];

src/query/expression/src/block.rs

@@ -23,6 +23,7 @@ use databend_common_exception::Result;
 
 use crate::schema::DataSchema;
 use crate::types::AnyType;
+use crate::types::ArgType;
 use crate::types::DataType;
 use crate::Column;
 use crate::ColumnBuilder;
@@ -71,6 +72,14 @@ impl BlockEntry {
         Self { data_type, value }
     }
 
+    pub fn from_arg_value<T: ArgType>(value: Value<T>) -> Self {
+        Self::new(T::data_type(), value.upcast())
+    }
+
+    pub fn from_arg_scalar<T: ArgType>(scalar: T::Scalar) -> Self {
+        Self::new(T::data_type(), Value::Scalar(T::upcast_scalar(scalar)))
+    }
+
     pub fn remove_nullable(self) -> Self {
         match self.value {
             Value::Column(Column::Nullable(col)) => {
@@ -89,6 +98,12 @@ impl BlockEntry {
     }
 }
 
+impl From<Column> for BlockEntry {
+    fn from(v: Column) -> Self {
+        Self::new(v.data_type(), Value::Column(v))
+    }
+}
+
 #[typetag::serde(tag = "type")]
 pub trait BlockMetaInfo: Debug + Send + Sync + Any + 'static {
     #[allow(clippy::borrowed_box)]

src/query/expression/src/evaluator.rs

@@ -2646,7 +2646,11 @@ impl ValueVisitor for CheckStrictValue {
         self.visit_column(column.column)
     }
 
-    fn visit_typed_column<T: ValueType>(&mut self, _: T::Column) -> std::result::Result<(), ()> {
+    fn visit_typed_column<T: ValueType>(
+        &mut self,
+        _: T::Column,
+        _: &DataType,
+    ) -> std::result::Result<(), ()> {
         Ok(())
     }
 }

src/query/expression/src/function.rs

@@ -320,7 +320,8 @@ impl Function {
         debug_assert!(!self.signature.return_type.is_nullable_or_null());
 
         let mut signature = self.signature;
-        signature.return_type = signature.return_type.wrap_nullable();
+        let return_type = signature.return_type;
+        signature.return_type = return_type.wrap_nullable();
 
         let (calc_domain, eval) = self.eval.into_scalar().unwrap();
 
@@ -332,7 +333,10 @@ impl Function {
                         has_null: false,
                         value: Some(Box::new(domain)),
                     };
-                    FunctionDomain::Domain(NullableType::<AnyType>::upcast_domain(new_domain))
+                    FunctionDomain::Domain(NullableType::<AnyType>::upcast_domain_with_type(
+                        new_domain,
+                        &return_type,
+                    ))
                 }
                 FunctionDomain::Full | FunctionDomain::MayThrow => FunctionDomain::Full,
             }
@@ -787,13 +791,13 @@ pub fn error_to_null<I1: AccessType, O: ArgType>(
             match output {
                 Value::Scalar(_) => Value::Scalar(None),
                 Value::Column(column) => {
-                    Value::Column(NullableColumn::new(column, validity.into()))
+                    Value::Column(NullableColumn::new_unchecked(column, validity.into()))
                 }
             }
         } else {
             match output {
                 Value::Scalar(scalar) => Value::Scalar(Some(scalar)),
-                Value::Column(column) => Value::Column(NullableColumn::new(
+                Value::Column(column) => Value::Column(NullableColumn::new_unchecked(
                     column,
                     Bitmap::new_constant(true, ctx.num_rows),
                 )),

src/query/expression/src/kernels/concat.rs

@@ -32,6 +32,7 @@ use crate::types::number::NumberColumn;
 use crate::types::vector::VectorColumnBuilder;
 use crate::types::AccessType;
 use crate::types::AnyType;
+use crate::types::ArgType;
 use crate::types::ArrayType;
 use crate::types::BooleanType;
 use crate::types::DataType;
@@ -128,6 +129,7 @@ impl Column {
             Some(col) => col,
         };
         let capacity = columns_iter_clone.fold(first_column.len(), |acc, x| acc + x.len());
+        let data_type = first_column.data_type();
         let column = match first_column {
             Column::Null { .. } => Column::Null { len: capacity },
             Column::EmptyArray { .. } => Column::EmptyArray { len: capacity },
@@ -182,7 +184,7 @@ impl Column {
                 offsets.push(0);
                 let builder = ColumnBuilder::with_capacity(&col.values().data_type(), capacity);
                 let builder = ArrayColumnBuilder { builder, offsets };
-                Self::concat_value_types::<ArrayType<AnyType>>(builder, columns)
+                Self::concat_value_types::<ArrayType<AnyType>>(builder, columns, &data_type)
             }
             Column::Map(col) => {
                 let mut offsets = Vec::with_capacity(capacity + 1);
@@ -199,7 +201,7 @@ impl Column {
                     values: val_builder,
                 };
                 let builder = ArrayColumnBuilder { builder, offsets };
-                Self::concat_value_types::<MapType<AnyType, AnyType>>(builder, columns)
+                Self::concat_value_types::<MapType<AnyType, AnyType>>(builder, columns, &data_type)
             }
             Column::Nullable(_) => {
                 let column: Vec<Column> = columns
@@ -291,11 +293,12 @@ impl Column {
     fn concat_value_types<T: ValueType>(
         mut builder: T::ColumnBuilder,
         columns: impl Iterator<Item = Column>,
+        data_type: &DataType,
     ) -> Column {
         let columns = columns.map(|c| T::try_downcast_column(&c).unwrap());
         for col in columns {
             T::append_column(&mut builder, &col);
         }
-        T::upcast_column(T::build_column(builder))
+        T::upcast_column_with_type(T::build_column(builder), data_type)
     }
 }

src/query/expression/src/kernels/filter.rs

@@ -203,8 +203,12 @@ impl ValueVisitor for FilterVisitor<'_> {
         Ok(())
     }
 
-    fn visit_typed_column<T: ValueType>(&mut self, column: T::Column) -> Result<()> {
-        let c = T::upcast_column(column.clone());
+    fn visit_typed_column<T: ValueType>(
+        &mut self,
+        column: T::Column,
+        data_type: &DataType,
+    ) -> Result<()> {
+        let c = T::upcast_column_with_type(column.clone(), data_type);
         let builder = ColumnBuilder::with_capacity(&c.data_type(), c.len());
         let mut builder = T::try_downcast_owned_builder(builder).unwrap();
         match self.strategy {
@@ -224,7 +228,10 @@ impl ValueVisitor for FilterVisitor<'_> {
             }
         }
 
-        self.result = Some(Value::Column(T::upcast_column(T::build_column(builder))));
+        self.result = Some(Value::Column(T::upcast_column_with_type(
+            T::build_column(builder),
+            data_type,
+        )));
         Ok(())
     }