diff --git a/src/query/expression/src/types.rs b/src/query/expression/src/types.rs index 80cb1ed209829..cbae779c3e18e 100755 --- a/src/query/expression/src/types.rs +++ b/src/query/expression/src/types.rs @@ -313,7 +313,7 @@ pub trait ValueType: Debug + Clone + PartialEq + Sized + 'static { builder: &'a mut ColumnBuilder, ) -> Option<&'a mut Self::ColumnBuilder>; - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option; + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option; fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option; diff --git a/src/query/expression/src/types/any.rs b/src/query/expression/src/types/any.rs index 55d3743c8c717..bc10f689ede14 100755 --- a/src/query/expression/src/types/any.rs +++ b/src/query/expression/src/types/any.rs @@ -64,7 +64,7 @@ impl ValueType for AnyType { Some(builder) } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { Some(builder) } diff --git a/src/query/expression/src/types/array.rs b/src/query/expression/src/types/array.rs index 0da341ae8a28f..f80f4d39a7cf4 100755 --- a/src/query/expression/src/types/array.rs +++ b/src/query/expression/src/types/array.rs @@ -81,10 +81,21 @@ impl ValueType for ArrayType { None } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + #[allow(clippy::manual_map)] + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { match builder { ColumnBuilder::Array(inner) => { let builder = T::try_downcast_owned_builder(inner.builder); + // ``` + // builder.map(|builder| ArrayColumnBuilder { + // builder, + // offsets: inner.offsets, + // }) + // ``` + // If we using the clippy recommend way like above, the compiler will complain: + // use of partially moved value: `inner`. + // That's rust borrow checker error, if we using the new borrow checker named polonius, + // everything goes fine, but polonius is very slow, so we allow manual map here. if let Some(builder) = builder { Some(ArrayColumnBuilder { builder, diff --git a/src/query/expression/src/types/bitmap.rs b/src/query/expression/src/types/bitmap.rs index 058522e5ce0aa..5b9788f9955b4 100644 --- a/src/query/expression/src/types/bitmap.rs +++ b/src/query/expression/src/types/bitmap.rs @@ -68,7 +68,7 @@ impl ValueType for BitmapType { } } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { match builder { ColumnBuilder::Bitmap(builder) => Some(builder), _ => None, diff --git a/src/query/expression/src/types/boolean.rs b/src/query/expression/src/types/boolean.rs index d9a00d0b65e05..4411ac495527d 100644 --- a/src/query/expression/src/types/boolean.rs +++ b/src/query/expression/src/types/boolean.rs @@ -75,7 +75,7 @@ impl ValueType for BooleanType { } } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { match builder { ColumnBuilder::Boolean(builder) => Some(builder), _ => None, diff --git a/src/query/expression/src/types/date.rs b/src/query/expression/src/types/date.rs index a705f7410e895..7975a861cda23 100644 --- a/src/query/expression/src/types/date.rs +++ b/src/query/expression/src/types/date.rs @@ -103,7 +103,7 @@ impl ValueType for DateType { } } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { match builder { ColumnBuilder::Date(builder) => Some(builder), _ => None, diff --git a/src/query/expression/src/types/decimal.rs b/src/query/expression/src/types/decimal.rs index fc2459ceb487a..23c4551efbe21 100644 --- a/src/query/expression/src/types/decimal.rs +++ b/src/query/expression/src/types/decimal.rs @@ -89,12 +89,12 @@ impl ValueType for DecimalType { Num::try_downcast_builder(builder) } - fn try_downcast_owned_builder<'a>(_builder: ColumnBuilder) -> Option { - todo!() + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { + Num::try_downcast_owned_builder(builder) } - fn try_upcast_column_builder(_builder: Self::ColumnBuilder) -> Option { - todo!() + fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option { + Num::try_upcast_column_builder(builder, Num::default_decimal_size()) } fn upcast_scalar(scalar: Self::Scalar) -> Scalar { @@ -299,6 +299,10 @@ pub trait Decimal: fn try_downcast_column(column: &Column) -> Option<(Buffer, DecimalSize)>; fn try_downcast_builder<'a>(builder: &'a mut ColumnBuilder) -> Option<&'a mut Vec>; + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option>; + + fn try_upcast_column_builder(builder: Vec, size: DecimalSize) -> Option; + fn try_downcast_scalar(scalar: &DecimalScalar) -> Option; fn try_downcast_domain(domain: &DecimalDomain) -> Option>; @@ -485,6 +489,19 @@ impl Decimal for i128 { } } + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option> { + match builder { + ColumnBuilder::Decimal(DecimalColumnBuilder::Decimal128(s, _)) => Some(s), + _ => None, + } + } + + fn try_upcast_column_builder(builder: Vec, size: DecimalSize) -> Option { + Some(ColumnBuilder::Decimal(DecimalColumnBuilder::Decimal128( + builder, size, + ))) + } + fn try_downcast_scalar<'a>(scalar: &DecimalScalar) -> Option { match scalar { DecimalScalar::Decimal128(val, _) => Some(*val), @@ -638,6 +655,19 @@ impl Decimal for i256 { } } + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option> { + match builder { + ColumnBuilder::Decimal(DecimalColumnBuilder::Decimal256(s, _)) => Some(s), + _ => None, + } + } + + fn try_upcast_column_builder(builder: Vec, size: DecimalSize) -> Option { + Some(ColumnBuilder::Decimal(DecimalColumnBuilder::Decimal256( + builder, size, + ))) + } + fn try_downcast_scalar<'a>(scalar: &DecimalScalar) -> Option { match scalar { DecimalScalar::Decimal256(val, _) => Some(*val), diff --git a/src/query/expression/src/types/empty_array.rs b/src/query/expression/src/types/empty_array.rs index 64c8421c01c5b..7207c5533d4e6 100644 --- a/src/query/expression/src/types/empty_array.rs +++ b/src/query/expression/src/types/empty_array.rs @@ -76,7 +76,7 @@ impl ValueType for EmptyArrayType { } } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { match builder { ColumnBuilder::EmptyArray { len } => Some(len), _ => None, diff --git a/src/query/expression/src/types/empty_map.rs b/src/query/expression/src/types/empty_map.rs index 253dd99d4f632..2176d221d8fc8 100644 --- a/src/query/expression/src/types/empty_map.rs +++ b/src/query/expression/src/types/empty_map.rs @@ -76,7 +76,7 @@ impl ValueType for EmptyMapType { } } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { match builder { ColumnBuilder::EmptyMap { len } => Some(len), _ => None, diff --git a/src/query/expression/src/types/generic.rs b/src/query/expression/src/types/generic.rs index 35a6b673d3f8d..b8ce249193994 100755 --- a/src/query/expression/src/types/generic.rs +++ b/src/query/expression/src/types/generic.rs @@ -67,7 +67,7 @@ impl ValueType for GenericType { Some(builder) } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { Some(builder) } diff --git a/src/query/expression/src/types/map.rs b/src/query/expression/src/types/map.rs index 9be355aa34e17..38d170e813f50 100755 --- a/src/query/expression/src/types/map.rs +++ b/src/query/expression/src/types/map.rs @@ -363,8 +363,8 @@ impl ValueType for MapType { as ValueType>::try_downcast_owned_builder(builder) } - fn try_upcast_column_builder(_builder: Self::ColumnBuilder) -> Option { - None + fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option { + as ValueType>::try_upcast_column_builder(builder) } 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 index b1097b76098a9..78c5737934c32 100644 --- a/src/query/expression/src/types/null.rs +++ b/src/query/expression/src/types/null.rs @@ -80,7 +80,7 @@ impl ValueType for NullType { } } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { match builder { ColumnBuilder::Null { len } => Some(len), _ => None, diff --git a/src/query/expression/src/types/nullable.rs b/src/query/expression/src/types/nullable.rs index d588c1658a6c1..78a267052bb7f 100755 --- a/src/query/expression/src/types/nullable.rs +++ b/src/query/expression/src/types/nullable.rs @@ -95,10 +95,21 @@ impl ValueType for NullableType { None } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + #[allow(clippy::manual_map)] + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { match builder { ColumnBuilder::Nullable(inner) => { let builder = T::try_downcast_owned_builder(inner.builder); + // ``` + // builder.map(|builder| NullableColumnBuilder { + // builder, + // validity: inner.validity, + // }) + // ``` + // If we using the clippy recommend way like above, the compiler will complain: + // use of partially moved value: `inner`. + // That's rust borrow checker error, if we using the new borrow checker named polonius, + // everything goes fine, but polonius is very slow, so we allow manual map here. if let Some(builder) = builder { Some(NullableColumnBuilder { builder, diff --git a/src/query/expression/src/types/number.rs b/src/query/expression/src/types/number.rs index 5f32b4239721d..4339036ea6261 100644 --- a/src/query/expression/src/types/number.rs +++ b/src/query/expression/src/types/number.rs @@ -136,7 +136,7 @@ impl ValueType for NumberType { } } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { match builder { ColumnBuilder::Number(num) => Num::try_downcast_owned_builder(num), _ => None, diff --git a/src/query/expression/src/types/string.rs b/src/query/expression/src/types/string.rs index fbf2c872e8415..ac88684b6a824 100644 --- a/src/query/expression/src/types/string.rs +++ b/src/query/expression/src/types/string.rs @@ -79,7 +79,7 @@ impl ValueType for StringType { } } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { match builder { ColumnBuilder::String(builder) => Some(builder), _ => None, diff --git a/src/query/expression/src/types/timestamp.rs b/src/query/expression/src/types/timestamp.rs index e6f7b9067aaeb..dde8687408e32 100644 --- a/src/query/expression/src/types/timestamp.rs +++ b/src/query/expression/src/types/timestamp.rs @@ -110,7 +110,7 @@ impl ValueType for TimestampType { } } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { match builder { ColumnBuilder::Timestamp(builder) => Some(builder), _ => None, diff --git a/src/query/expression/src/types/variant.rs b/src/query/expression/src/types/variant.rs index 61398448d551c..99ba6f8a30835 100644 --- a/src/query/expression/src/types/variant.rs +++ b/src/query/expression/src/types/variant.rs @@ -87,7 +87,7 @@ impl ValueType for VariantType { } } - fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option { + fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option { match builder { ColumnBuilder::Variant(builder) => Some(builder), _ => None, diff --git a/src/query/functions/src/aggregates/aggregate_unary.rs b/src/query/functions/src/aggregates/aggregate_unary.rs index 0002d0939a7d4..f30384de5cebf 100644 --- a/src/query/functions/src/aggregates/aggregate_unary.rs +++ b/src/query/functions/src/aggregates/aggregate_unary.rs @@ -30,6 +30,8 @@ use common_expression::ColumnBuilder; use common_expression::Scalar; use common_expression::StateAddr; +use crate::aggregates::take_mut; + pub trait UnaryState: Send + Sync + Default where T: ValueType, @@ -226,8 +228,14 @@ where fn merge_result(&self, place: StateAddr, builder: &mut ColumnBuilder) -> Result<()> { let state: &mut S = place.get::(); - let builder = R::try_downcast_builder(builder).unwrap(); - state.merge_result(builder, self.function_data.as_deref())?; + // some `ValueType` like `NullableType` need ownership to downcast builder, + // so here we using an unsafe way to take the ownership of builder. + take_mut(builder, |builder| { + let mut builder = R::try_downcast_owned_builder(builder).unwrap(); + state + .merge_result(&mut builder, self.function_data.as_deref()) + .map(|_| R::try_upcast_column_builder(builder).unwrap()) + })?; Ok(()) } diff --git a/src/query/functions/src/aggregates/aggregator_common.rs b/src/query/functions/src/aggregates/aggregator_common.rs index 66e255fef4c2d..210b56f4cc8a1 100644 --- a/src/query/functions/src/aggregates/aggregator_common.rs +++ b/src/query/functions/src/aggregates/aggregator_common.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::fmt::Display; +use std::panic; use bumpalo::Bump; use common_exception::ErrorCode; @@ -167,3 +168,24 @@ pub fn deserialize_state(slice: &mut &[u8]) -> R *slice = &slice[bytes_read..]; Ok(value) } + +// copy from https://docs.rs/take_mut/0.2.2/take_mut/fn.take.html with some modifications. +// if a panic occurs, the entire process will be aborted, as there's no valid `T` to put back into the `&mut T`. +pub fn take_mut(mut_ref: &mut T, closure: F) -> Result<()> +where F: FnOnce(T) -> Result { + use std::ptr; + + unsafe { + let old_t = ptr::read(mut_ref); + let closure_result = panic::catch_unwind(panic::AssertUnwindSafe(|| closure(old_t))); + + match closure_result { + Ok(Ok(new_t)) => { + ptr::write(mut_ref, new_t); + Ok(()) + } + Ok(Err(e)) => Err(e), + Err(_) => ::std::process::abort(), + } + } +}