chore(query): refactoring ValueType, introducing ColumnBuilderMut (#18102)

Author: forsaken628

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

@@ -209,29 +209,24 @@ impl ValueVisitor for FilterVisitor<'_> {
         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();
+        let mut builder = ColumnBuilder::with_capacity(&c.data_type(), c.len());
+        let mut inner_builder = T::downcast_builder(&mut builder);
         match self.strategy {
             IterationStrategy::IndexIterator => {
                 let iter = TrueIdxIter::new(self.original_rows, Some(self.filter));
                 iter.for_each(|index| {
-                    T::push_item(&mut builder, unsafe {
-                        T::index_column_unchecked(&column, index)
-                    })
+                    inner_builder.push_item(unsafe { T::index_column_unchecked(&column, index) });
                 });
             }
             _ => {
                 let iter = SlicesIterator::new(self.filter);
                 iter.for_each(|(start, len)| {
-                    T::append_column(&mut builder, &T::slice_column(&column, start..start + len))
+                    inner_builder.append_column(&T::slice_column(&column, start..start + len));
                 });
             }
         }
-
-        self.result = Some(Value::Column(T::upcast_column_with_type(
-            T::build_column(builder),
-            data_type,
-        )));
+        drop(inner_builder);
+        self.result = Some(Value::Column(builder.build()));
         Ok(())
     }
 

src/query/expression/src/kernels/take.rs

@@ -160,18 +160,15 @@ where I: databend_common_column::types::Index
         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();
+        let mut builder = ColumnBuilder::with_capacity(&c.data_type(), c.len());
+        let mut inner_builder = T::downcast_builder(&mut builder);
 
         for index in self.indices {
-            T::push_item(&mut builder, unsafe {
-                T::index_column_unchecked(&column, index.to_usize())
-            });
+            inner_builder
+                .push_item(unsafe { T::index_column_unchecked(&column, index.to_usize()) });
         }
-        self.result = Some(Value::Column(T::upcast_column_with_type(
-            T::build_column(builder),
-            data_type,
-        )));
+        drop(inner_builder);
+        self.result = Some(Value::Column(builder.build()));
         Ok(())
     }
 

src/query/expression/src/kernels/take_compact.rs

@@ -106,21 +106,17 @@ impl ValueVisitor for TakeCompactVisitor<'_> {
         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();
+        let mut builder = ColumnBuilder::with_capacity(&c.data_type(), c.len());
+        let mut inner_builder = T::downcast_builder(&mut builder);
 
         for (index, cnt) in self.indices {
             for _ in 0..*cnt {
-                T::push_item(&mut builder, unsafe {
-                    T::index_column_unchecked(&column, *index as usize)
-                });
+                inner_builder
+                    .push_item(unsafe { T::index_column_unchecked(&column, *index as usize) });
             }
         }
-        self.result = Some(Value::Column(T::upcast_column_with_type(
-            T::build_column(builder),
-            data_type,
-        )));
-
+        drop(inner_builder);
+        self.result = Some(Value::Column(builder.build()));
         Ok(())
     }
 

src/query/expression/src/kernels/take_ranges.rs

@@ -108,21 +108,16 @@ impl ValueVisitor for TakeRangeVisitor<'_> {
         column: T::Column,
         data_type: &DataType,
     ) -> Result<()> {
-        let builder = ColumnBuilder::with_capacity(data_type, T::column_len(&column));
-        let mut builder = T::try_downcast_owned_builder(builder).unwrap();
+        let mut builder = ColumnBuilder::with_capacity(data_type, T::column_len(&column));
+        let mut inner_builder = T::downcast_builder(&mut builder);
 
         for range in self.ranges {
             for index in range.start as usize..range.end as usize {
-                T::push_item(&mut builder, unsafe {
-                    T::index_column_unchecked(&column, index)
-                });
+                inner_builder.push_item(unsafe { T::index_column_unchecked(&column, index) });
             }
         }
-        self.result = Some(Value::Column(T::upcast_column_with_type(
-            T::build_column(builder),
-            data_type,
-        )));
-
+        drop(inner_builder);
+        self.result = Some(Value::Column(builder.build()));
         Ok(())
     }
 

src/query/expression/src/lib.rs

@@ -42,6 +42,7 @@
 #![feature(alloc_layout_extra)]
 #![feature(debug_closure_helpers)]
 #![feature(never_type)]
+#![feature(slice_as_array)]
 
 #[allow(dead_code)]
 mod block;

src/query/expression/src/types.rs

@@ -40,6 +40,8 @@ pub mod zero_size_type;
 use std::cmp::Ordering;
 use std::fmt::Debug;
 use std::iter::TrustedLen;
+use std::ops::Deref;
+use std::ops::DerefMut;
 use std::ops::Range;
 
 use borsh::BorshDeserialize;
@@ -456,6 +458,7 @@ pub trait AccessType: Debug + Clone + PartialEq + Sized + 'static {
 /// [ValueType] includes the builder method of a data type based on [AccessType].
 pub trait ValueType: AccessType {
     type ColumnBuilder: Debug + Clone;
+    type ColumnBuilderMut<'a>: Debug + From<&'a mut Self::ColumnBuilder> + BuilderExt<Self>;
 
     /// Convert a scalar value to the generic [Scalar] type
     fn upcast_scalar_with_type(scalar: Self::Scalar, data_type: &DataType) -> Scalar;
@@ -466,27 +469,7 @@ pub trait ValueType: AccessType {
     /// Convert a column value to the generic Column type
     fn upcast_column_with_type(col: Self::Column, data_type: &DataType) -> Column;
 
-    /// Downcast `ColumnBuilder` to a mutable reference of its inner builder type.
-    ///
-    /// Not every builder can be downcasted successfully.
-    /// For example: `ArrayType<T: ValueType>`, `NullableType<T: ValueType>`, and `KvPair<K: ValueType, V: ValueType>`
-    /// cannot be downcasted and this method will return `None`.
-    ///
-    /// So when using this method, we cannot unwrap the returned value directly.
-    /// We should:
-    ///
-    /// ```ignore
-    /// // builder: ColumnBuilder
-    /// // T: ValueType
-    /// if let Some(inner) = T::try_downcast_builder(&mut builder) {
-    ///     inner.push(...);
-    /// } else {
-    ///     builder.push(...);
-    /// }
-    /// ```
-    fn try_downcast_builder(builder: &mut ColumnBuilder) -> Option<&mut Self::ColumnBuilder>;
-
-    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder>;
+    fn downcast_builder(builder: &mut ColumnBuilder) -> Self::ColumnBuilderMut<'_>;
 
     fn try_upcast_column_builder(
         builder: Self::ColumnBuilder,
@@ -495,15 +478,92 @@ pub trait ValueType: AccessType {
 
     fn column_to_builder(col: Self::Column) -> Self::ColumnBuilder;
 
+    fn builder_len_mut(builder: &Self::ColumnBuilderMut<'_>) -> usize;
     fn builder_len(builder: &Self::ColumnBuilder) -> usize;
-    fn push_item(builder: &mut Self::ColumnBuilder, item: Self::ScalarRef<'_>);
-    fn push_item_repeat(builder: &mut Self::ColumnBuilder, item: Self::ScalarRef<'_>, n: usize);
-    fn push_default(builder: &mut Self::ColumnBuilder);
-    fn append_column(builder: &mut Self::ColumnBuilder, other: &Self::Column);
+
+    fn push_item_mut(builder: &mut Self::ColumnBuilderMut<'_>, item: Self::ScalarRef<'_>);
+    fn push_item(builder: &mut Self::ColumnBuilder, item: Self::ScalarRef<'_>) {
+        Self::push_item_mut(&mut builder.into(), item);
+    }
+
+    fn push_item_repeat_mut(
+        builder: &mut Self::ColumnBuilderMut<'_>,
+        item: Self::ScalarRef<'_>,
+        n: usize,
+    );
+    fn push_item_repeat(builder: &mut Self::ColumnBuilder, item: Self::ScalarRef<'_>, n: usize) {
+        Self::push_item_repeat_mut(&mut builder.into(), item, n);
+    }
+
+    fn push_default_mut(builder: &mut Self::ColumnBuilderMut<'_>);
+    fn push_default(builder: &mut Self::ColumnBuilder) {
+        Self::push_default_mut(&mut builder.into())
+    }
+
+    fn append_column_mut(builder: &mut Self::ColumnBuilderMut<'_>, other: &Self::Column);
+    fn append_column(builder: &mut Self::ColumnBuilder, other: &Self::Column) {
+        Self::append_column_mut(&mut builder.into(), other);
+    }
+
     fn build_column(builder: Self::ColumnBuilder) -> Self::Column;
     fn build_scalar(builder: Self::ColumnBuilder) -> Self::Scalar;
 }
 
+impl<'a, T: ValueType> From<&'a mut T::ColumnBuilder> for BuilderMut<'a, T> {
+    fn from(value: &'a mut T::ColumnBuilder) -> Self {
+        BuilderMut(value)
+    }
+}
+
+#[derive(Debug)]
+pub struct BuilderMut<'a, T: ValueType>(&'a mut T::ColumnBuilder);
+
+impl<T: ValueType> Deref for BuilderMut<'_, T> {
+    type Target = T::ColumnBuilder;
+
+    fn deref(&self) -> &Self::Target {
+        self.0
+    }
+}
+
+impl<T: ValueType> DerefMut for BuilderMut<'_, T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.0
+    }
+}
+
+impl<'a, T> BuilderExt<T> for BuilderMut<'a, T>
+where T: ValueType<ColumnBuilderMut<'a> = Self>
+{
+    fn len(&self) -> usize {
+        T::builder_len(self.0)
+    }
+
+    fn push_item(&mut self, item: <T>::ScalarRef<'_>) {
+        T::push_item_mut(self, item);
+    }
+
+    fn push_repeat(&mut self, item: <T>::ScalarRef<'_>, n: usize) {
+        T::push_item_repeat_mut(self, item, n);
+    }
+
+    fn push_default(&mut self) {
+        T::push_default_mut(self);
+    }
+
+    fn append_column(&mut self, other: &<T>::Column) {
+        T::append_column_mut(self, other);
+    }
+}
+
+pub trait BuilderExt<T: ValueType> {
+    fn len(&self) -> usize;
+    fn push_item(&mut self, item: T::ScalarRef<'_>);
+    fn push_repeat(&mut self, item: T::ScalarRef<'_>, n: usize);
+    fn push_default(&mut self);
+    fn append_column(&mut self, other: &T::Column);
+}
+
 /// Almost all [ValueType] implement [ReturnType], except [AnyType].
 pub trait ReturnType: ValueType {
     fn create_builder(capacity: usize, generics: &GenericMap) -> Self::ColumnBuilder;

src/query/expression/src/types/any.rs

@@ -17,6 +17,7 @@ use std::ops::Range;
 
 use crate::property::Domain;
 use crate::types::AccessType;
+use crate::types::BuilderMut;
 use crate::types::DataType;
 use crate::types::ValueType;
 use crate::values::Column;
@@ -92,6 +93,7 @@ impl AccessType for AnyType {
 
 impl ValueType for AnyType {
     type ColumnBuilder = ColumnBuilder;
+    type ColumnBuilderMut<'a> = BuilderMut<'a, Self>;
 
     fn upcast_scalar_with_type(scalar: Self::Scalar, _: &DataType) -> Scalar {
         scalar
@@ -106,12 +108,8 @@ impl ValueType for AnyType {
         col
     }
 
-    fn try_downcast_builder(builder: &mut ColumnBuilder) -> Option<&mut Self::ColumnBuilder> {
-        Some(builder)
-    }
-
-    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder> {
-        Some(builder)
+    fn downcast_builder(builder: &mut ColumnBuilder) -> Self::ColumnBuilderMut<'_> {
+        builder.into()
     }
 
     fn try_upcast_column_builder(
@@ -129,19 +127,27 @@ impl ValueType for AnyType {
         builder.len()
     }
 
-    fn push_item(builder: &mut Self::ColumnBuilder, item: Self::ScalarRef<'_>) {
+    fn builder_len_mut(builder: &Self::ColumnBuilderMut<'_>) -> usize {
+        builder.len()
+    }
+
+    fn push_item_mut(builder: &mut Self::ColumnBuilderMut<'_>, item: Self::ScalarRef<'_>) {
         builder.push(item);
     }
 
-    fn push_item_repeat(builder: &mut Self::ColumnBuilder, item: Self::ScalarRef<'_>, n: usize) {
+    fn push_item_repeat_mut(
+        builder: &mut Self::ColumnBuilderMut<'_>,
+        item: Self::ScalarRef<'_>,
+        n: usize,
+    ) {
         builder.push_repeat(&item, n);
     }
 
-    fn push_default(builder: &mut Self::ColumnBuilder) {
+    fn push_default_mut(builder: &mut Self::ColumnBuilderMut<'_>) {
         builder.push_default();
     }
 
-    fn append_column(builder: &mut Self::ColumnBuilder, other: &Self::Column) {
+    fn append_column_mut(builder: &mut Self::ColumnBuilderMut<'_>, other: &Self::Column) {
         builder.append_column(other);
     }