support decimal for quantile_cont

databendlabs · Feb 25, 2024 · 95b0eff · 95b0eff
1 parent ffccc1f
commit 95b0eff
Show file tree

Hide file tree

Showing 22 changed files with 379 additions and 61 deletions.
diff --git a/src/query/expression/src/types.rs b/src/query/expression/src/types.rs
@@ -345,7 +345,10 @@ pub trait ValueType: Debug + Clone + PartialEq + Sized + 'static {
 
     fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder>;
 
-    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder>;
+    fn try_upcast_column_builder(
+        builder: Self::ColumnBuilder,
+        decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder>;
 
     fn upcast_scalar(scalar: Self::Scalar) -> Scalar;
     fn upcast_column(col: Self::Column) -> Column;

diff --git a/src/query/expression/src/types/any.rs b/src/query/expression/src/types/any.rs
@@ -16,6 +16,7 @@ use std::cmp::Ordering;
 use std::ops::Range;
 
 use crate::property::Domain;
+use crate::types::DecimalSize;
 use crate::types::ValueType;
 use crate::values::Column;
 use crate::values::Scalar;
@@ -67,7 +68,10 @@ impl ValueType for AnyType {
         Some(builder)
     }
 
-    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+    fn try_upcast_column_builder(
+        builder: Self::ColumnBuilder,
+        _decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
         Some(builder)
     }
 

diff --git a/src/query/expression/src/types/array.rs b/src/query/expression/src/types/array.rs
@@ -22,6 +22,7 @@ use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 
 use super::AnyType;
+use super::DecimalSize;
 use crate::property::Domain;
 use crate::types::ArgType;
 use crate::types::DataType;
@@ -107,8 +108,11 @@ impl<T: ValueType> ValueType for ArrayType<T> {
         }
     }
 
-    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
-        Some(ColumnBuilder::Array(Box::new(builder.upcast())))
+    fn try_upcast_column_builder(
+        builder: Self::ColumnBuilder,
+        decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
+        Some(ColumnBuilder::Array(Box::new(builder.upcast(decimal_size))))
     }
 
     fn upcast_scalar(scalar: Self::Scalar) -> Scalar {
@@ -395,9 +399,9 @@ impl<T: ValueType> ArrayColumnBuilder<T> {
         )
     }
 
-    pub fn upcast(self) -> ArrayColumnBuilder<AnyType> {
+    pub fn upcast(self, decimal_size: Option<DecimalSize>) -> ArrayColumnBuilder<AnyType> {
         ArrayColumnBuilder {
-            builder: T::try_upcast_column_builder(self.builder).unwrap(),
+            builder: T::try_upcast_column_builder(self.builder, decimal_size).unwrap(),
             offsets: self.offsets,
         }
     }

diff --git a/src/query/expression/src/types/binary.rs b/src/query/expression/src/types/binary.rs
@@ -25,6 +25,7 @@ use serde::Serialize;
 use crate::property::Domain;
 use crate::types::ArgType;
 use crate::types::DataType;
+use crate::types::DecimalSize;
 use crate::types::GenericMap;
 use crate::types::ValueType;
 use crate::utils::arrow::buffer_into_mut;
@@ -87,7 +88,10 @@ impl ValueType for BinaryType {
         }
     }
 
-    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+    fn try_upcast_column_builder(
+        builder: Self::ColumnBuilder,
+        _decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
         Some(ColumnBuilder::Binary(builder))
     }
 

diff --git a/src/query/expression/src/types/bitmap.rs b/src/query/expression/src/types/bitmap.rs
@@ -20,6 +20,7 @@ use super::binary::BinaryIterator;
 use crate::property::Domain;
 use crate::types::ArgType;
 use crate::types::DataType;
+use crate::types::DecimalSize;
 use crate::types::GenericMap;
 use crate::types::ValueType;
 use crate::values::Column;
@@ -73,7 +74,10 @@ impl ValueType for BitmapType {
         }
     }
 
-    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+    fn try_upcast_column_builder(
+        builder: Self::ColumnBuilder,
+        _decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
         Some(ColumnBuilder::Bitmap(builder))
     }
 

diff --git a/src/query/expression/src/types/boolean.rs b/src/query/expression/src/types/boolean.rs
@@ -20,6 +20,7 @@ use databend_common_arrow::arrow::bitmap::MutableBitmap;
 use crate::property::Domain;
 use crate::types::ArgType;
 use crate::types::DataType;
+use crate::types::DecimalSize;
 use crate::types::GenericMap;
 use crate::types::ValueType;
 use crate::utils::arrow::bitmap_into_mut;
@@ -80,7 +81,10 @@ impl ValueType for BooleanType {
         }
     }
 
-    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+    fn try_upcast_column_builder(
+        builder: Self::ColumnBuilder,
+        _decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
         Some(ColumnBuilder::Boolean(builder))
     }
 

diff --git a/src/query/expression/src/types/date.rs b/src/query/expression/src/types/date.rs
@@ -28,6 +28,7 @@ use crate::date_helper::DateConverter;
 use crate::property::Domain;
 use crate::types::ArgType;
 use crate::types::DataType;
+use crate::types::DecimalSize;
 use crate::types::GenericMap;
 use crate::types::ValueType;
 use crate::utils::arrow::buffer_into_mut;
@@ -108,7 +109,10 @@ impl ValueType for DateType {
         }
     }
 
-    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+    fn try_upcast_column_builder(
+        builder: Self::ColumnBuilder,
+        _decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
         Some(ColumnBuilder::Date(builder))
     }
 

diff --git a/src/query/expression/src/types/decimal.rs b/src/query/expression/src/types/decimal.rs
@@ -94,8 +94,14 @@ impl<Num: Decimal> ValueType for DecimalType<Num> {
         Num::try_downcast_owned_builder(builder)
     }
 
-    fn try_upcast_column_builder(_builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
-        None
+    fn try_upcast_column_builder(
+        builder: Self::ColumnBuilder,
+        decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
+        Some(ColumnBuilder::Decimal(Num::upcast_builder(
+            builder,
+            decimal_size.unwrap(),
+        )))
     }
 
     fn upcast_scalar(scalar: Self::Scalar) -> Scalar {
@@ -360,6 +366,7 @@ pub trait Decimal:
     fn upcast_scalar(scalar: Self, size: DecimalSize) -> Scalar;
     fn upcast_column(col: Buffer<Self>, size: DecimalSize) -> Column;
     fn upcast_domain(domain: SimpleDomain<Self>, size: DecimalSize) -> Domain;
+    fn upcast_builder(builder: Vec<Self>, size: DecimalSize) -> DecimalColumnBuilder;
     fn data_type() -> DataType;
     const MIN: Self;
     const MAX: Self;
@@ -574,6 +581,10 @@ impl Decimal for i128 {
         Domain::Decimal(DecimalDomain::Decimal128(domain, size))
     }
 
+    fn upcast_builder(builder: Vec<Self>, size: DecimalSize) -> DecimalColumnBuilder {
+        DecimalColumnBuilder::Decimal128(builder, size)
+    }
+
     fn data_type() -> DataType {
         DataType::Decimal(DecimalDataType::Decimal128(DecimalSize {
             precision: MAX_DECIMAL128_PRECISION,
@@ -733,6 +744,10 @@ impl Decimal for i256 {
         Domain::Decimal(DecimalDomain::Decimal256(domain, size))
     }
 
+    fn upcast_builder(builder: Vec<Self>, size: DecimalSize) -> DecimalColumnBuilder {
+        DecimalColumnBuilder::Decimal256(builder, size)
+    }
+
     fn data_type() -> DataType {
         DataType::Decimal(DecimalDataType::Decimal256(DecimalSize {
             precision: MAX_DECIMAL256_PRECISION,
@@ -1080,6 +1095,13 @@ impl DecimalColumnBuilder {
             }
         })
     }
+
+    pub fn decimal_size(&self) -> DecimalSize {
+        match self {
+            DecimalColumnBuilder::Decimal128(_, size)
+            | DecimalColumnBuilder::Decimal256(_, size) => *size,
+        }
+    }
 }
 
 impl PartialOrd for DecimalScalar {

diff --git a/src/query/expression/src/types/empty_array.rs b/src/query/expression/src/types/empty_array.rs
@@ -17,6 +17,7 @@ use std::ops::Range;
 use crate::property::Domain;
 use crate::types::ArgType;
 use crate::types::DataType;
+use crate::types::DecimalSize;
 use crate::types::GenericMap;
 use crate::types::ValueType;
 use crate::values::Column;
@@ -81,7 +82,10 @@ impl ValueType for EmptyArrayType {
         }
     }
 
-    fn try_upcast_column_builder(len: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+    fn try_upcast_column_builder(
+        len: Self::ColumnBuilder,
+        _decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
         Some(ColumnBuilder::EmptyArray { len })
     }
 

diff --git a/src/query/expression/src/types/empty_map.rs b/src/query/expression/src/types/empty_map.rs
@@ -17,6 +17,7 @@ use std::ops::Range;
 use crate::property::Domain;
 use crate::types::ArgType;
 use crate::types::DataType;
+use crate::types::DecimalSize;
 use crate::types::GenericMap;
 use crate::types::ValueType;
 use crate::values::Column;
@@ -81,7 +82,10 @@ impl ValueType for EmptyMapType {
         }
     }
 
-    fn try_upcast_column_builder(len: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+    fn try_upcast_column_builder(
+        len: Self::ColumnBuilder,
+        _decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
         Some(ColumnBuilder::EmptyMap { len })
     }
 

diff --git a/src/query/expression/src/types/generic.rs b/src/query/expression/src/types/generic.rs
@@ -17,6 +17,7 @@ use std::ops::Range;
 use crate::property::Domain;
 use crate::types::ArgType;
 use crate::types::DataType;
+use crate::types::DecimalSize;
 use crate::types::GenericMap;
 use crate::types::ValueType;
 use crate::values::Column;
@@ -69,7 +70,10 @@ impl<const INDEX: usize> ValueType for GenericType<INDEX> {
         Some(builder)
     }
 
-    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+    fn try_upcast_column_builder(
+        builder: Self::ColumnBuilder,
+        _decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
         Some(builder)
     }
 

diff --git a/src/query/expression/src/types/geometry.rs b/src/query/expression/src/types/geometry.rs
@@ -25,6 +25,7 @@ use super::binary::BinaryIterator;
 use crate::property::Domain;
 use crate::types::ArgType;
 use crate::types::DataType;
+use crate::types::DecimalSize;
 use crate::types::GenericMap;
 use crate::types::ValueType;
 use crate::values::Column;
@@ -86,7 +87,10 @@ impl ValueType for GeometryType {
         }
     }
 
-    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+    fn try_upcast_column_builder(
+        builder: Self::ColumnBuilder,
+        _decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
         Some(ColumnBuilder::Geometry(builder))
     }
 

diff --git a/src/query/expression/src/types/map.rs b/src/query/expression/src/types/map.rs
@@ -18,6 +18,7 @@ use std::ops::Range;
 use databend_common_arrow::arrow::trusted_len::TrustedLen;
 
 use super::ArrayType;
+use super::DecimalSize;
 use crate::property::Domain;
 use crate::types::array::ArrayColumn;
 use crate::types::ArgType;
@@ -92,7 +93,10 @@ impl<K: ValueType, V: ValueType> ValueType for KvPair<K, V> {
         None
     }
 
-    fn try_upcast_column_builder(_builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+    fn try_upcast_column_builder(
+        _builder: Self::ColumnBuilder,
+        _decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
         None
     }
 
@@ -358,8 +362,11 @@ impl<K: ValueType, V: ValueType> ValueType for MapType<K, V> {
         <MapInternal<K, V> as ValueType>::try_downcast_owned_builder(builder)
     }
 
-    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
-        <MapInternal<K, V> as ValueType>::try_upcast_column_builder(builder)
+    fn try_upcast_column_builder(
+        builder: Self::ColumnBuilder,
+        decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
+        <MapInternal<K, V> as ValueType>::try_upcast_column_builder(builder, decimal_size)
     }
 
     fn upcast_scalar(scalar: Self::Scalar) -> Scalar {

diff --git a/src/query/expression/src/types/null.rs b/src/query/expression/src/types/null.rs
@@ -18,6 +18,7 @@ use super::nullable::NullableDomain;
 use crate::property::Domain;
 use crate::types::ArgType;
 use crate::types::DataType;
+use crate::types::DecimalSize;
 use crate::types::GenericMap;
 use crate::types::ValueType;
 use crate::values::Column;
@@ -85,7 +86,10 @@ impl ValueType for NullType {
         }
     }
 
-    fn try_upcast_column_builder(len: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+    fn try_upcast_column_builder(
+        len: Self::ColumnBuilder,
+        _decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
         Some(ColumnBuilder::Null { len })
     }
 

diff --git a/src/query/expression/src/types/nullable.rs b/src/query/expression/src/types/nullable.rs
@@ -20,6 +20,7 @@ use databend_common_arrow::arrow::bitmap::MutableBitmap;
 use databend_common_arrow::arrow::trusted_len::TrustedLen;
 
 use super::AnyType;
+use super::DecimalSize;
 use crate::property::Domain;
 use crate::types::ArgType;
 use crate::types::DataType;
@@ -121,8 +122,13 @@ impl<T: ValueType> ValueType for NullableType<T> {
         }
     }
 
-    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
-        Some(ColumnBuilder::Nullable(Box::new(builder.upcast())))
+    fn try_upcast_column_builder(
+        builder: Self::ColumnBuilder,
+        decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
+        Some(ColumnBuilder::Nullable(Box::new(
+            builder.upcast(decimal_size),
+        )))
     }
 
     fn upcast_scalar(scalar: Self::Scalar) -> Scalar {
@@ -377,9 +383,9 @@ impl<T: ValueType> NullableColumnBuilder<T> {
         }
     }
 
-    pub fn upcast(self) -> NullableColumnBuilder<AnyType> {
+    pub fn upcast(self, decimal_size: Option<DecimalSize>) -> NullableColumnBuilder<AnyType> {
         NullableColumnBuilder {
-            builder: T::try_upcast_column_builder(self.builder).unwrap(),
+            builder: T::try_upcast_column_builder(self.builder, decimal_size).unwrap(),
             validity: self.validity,
         }
     }

diff --git a/src/query/expression/src/types/number.rs b/src/query/expression/src/types/number.rs
@@ -142,7 +142,10 @@ impl<Num: Number> ValueType for NumberType<Num> {
         }
     }
 
-    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+    fn try_upcast_column_builder(
+        builder: Self::ColumnBuilder,
+        _decimal_size: Option<DecimalSize>,
+    ) -> Option<ColumnBuilder> {
         Num::try_upcast_column_builder(builder)
     }