add downcast_owned_builder

databendlabs · Nov 22, 2023 · 217bcc4 · 217bcc4
1 parent ea9091d
commit 217bcc4
Show file tree

Hide file tree

Showing 22 changed files with 447 additions and 15 deletions.
diff --git a/src/common/base/src/base/mod.rs b/src/common/base/src/base/mod.rs
@@ -22,6 +22,7 @@ mod singleton_instance;
 mod stop_handle;
 mod stoppable;
 mod string;
+mod take_mut;
 mod uniq_id;
 
 pub use net::get_free_tcp_port;
@@ -46,6 +47,7 @@ pub use string::mask_connection_info;
 pub use string::mask_string;
 pub use string::unescape_for_key;
 pub use string::unescape_string;
+pub use take_mut::take_mut;
 pub use tokio;
 pub use uniq_id::GlobalSequence;
 pub use uniq_id::GlobalUniqName;

diff --git a/src/common/base/src/base/take_mut.rs b/src/common/base/src/base/take_mut.rs
@@ -0,0 +1,38 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::panic;
+
+use common_exception::Result;
+
+/// 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<T, F>(mut_ref: &mut T, closure: F) -> Result<()>
+where F: FnOnce(T) -> Result<T> {
+    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(),
+        }
+    }
+}
diff --git a/src/query/expression/src/types.rs b/src/query/expression/src/types.rs
@@ -313,6 +313,10 @@ pub trait ValueType: Debug + Clone + PartialEq + Sized + 'static {
         builder: &'a mut ColumnBuilder,
     ) -> Option<&'a mut Self::ColumnBuilder>;
 
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder>;
+
+    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder>;
+
     fn upcast_scalar(scalar: Self::Scalar) -> Scalar;
     fn upcast_column(col: Self::Column) -> Column;
     fn upcast_domain(domain: Self::Domain) -> Domain;

diff --git a/src/query/expression/src/types/any.rs b/src/query/expression/src/types/any.rs
@@ -64,6 +64,14 @@ impl ValueType for AnyType {
         Some(builder)
     }
 
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder> {
+        Some(builder)
+    }
+
+    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+        Some(builder)
+    }
+
     fn upcast_scalar(scalar: Self::Scalar) -> Scalar {
         scalar
     }

diff --git a/src/query/expression/src/types/array.rs b/src/query/expression/src/types/array.rs
@@ -81,6 +81,38 @@ impl<T: ValueType> ValueType for ArrayType<T> {
         None
     }
 
+    #[allow(clippy::manual_map)]
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder> {
+        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,
+                        offsets: inner.offsets,
+                    })
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        }
+    }
+
+    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+        Some(ColumnBuilder::Array(Box::new(builder.upcast())))
+    }
+
     fn upcast_scalar(scalar: Self::Scalar) -> Scalar {
         Scalar::Array(T::upcast_column(scalar))
     }
@@ -366,6 +398,13 @@ impl<T: ValueType> ArrayColumnBuilder<T> {
             (self.offsets[0] as usize)..(self.offsets[1] as usize),
         )
     }
+
+    pub fn upcast(self) -> ArrayColumnBuilder<AnyType> {
+        ArrayColumnBuilder {
+            builder: T::try_upcast_column_builder(self.builder).unwrap(),
+            offsets: self.offsets,
+        }
+    }
 }
 
 impl<T: ArgType> ArrayColumnBuilder<T> {

diff --git a/src/query/expression/src/types/bitmap.rs b/src/query/expression/src/types/bitmap.rs
@@ -63,11 +63,22 @@ impl ValueType for BitmapType {
         builder: &'a mut ColumnBuilder,
     ) -> Option<&'a mut Self::ColumnBuilder> {
         match builder {
-            crate::ColumnBuilder::Bitmap(builder) => Some(builder),
+            ColumnBuilder::Bitmap(builder) => Some(builder),
             _ => None,
         }
     }
 
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder> {
+        match builder {
+            ColumnBuilder::Bitmap(builder) => Some(builder),
+            _ => None,
+        }
+    }
+
+    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+        Some(ColumnBuilder::Bitmap(builder))
+    }
+
     fn try_downcast_domain(domain: &Domain) -> Option<Self::Domain> {
         if domain.is_undefined() {
             Some(())

diff --git a/src/query/expression/src/types/boolean.rs b/src/query/expression/src/types/boolean.rs
@@ -70,11 +70,22 @@ impl ValueType for BooleanType {
         builder: &'a mut ColumnBuilder,
     ) -> Option<&'a mut Self::ColumnBuilder> {
         match builder {
-            crate::ColumnBuilder::Boolean(builder) => Some(builder),
+            ColumnBuilder::Boolean(builder) => Some(builder),
             _ => None,
         }
     }
 
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder> {
+        match builder {
+            ColumnBuilder::Boolean(builder) => Some(builder),
+            _ => None,
+        }
+    }
+
+    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+        Some(ColumnBuilder::Boolean(builder))
+    }
+
     fn try_downcast_domain(domain: &Domain) -> Option<Self::Domain> {
         domain.as_boolean().map(BooleanDomain::clone)
     }

diff --git a/src/query/expression/src/types/date.rs b/src/query/expression/src/types/date.rs
@@ -103,6 +103,17 @@ impl ValueType for DateType {
         }
     }
 
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder> {
+        match builder {
+            ColumnBuilder::Date(builder) => Some(builder),
+            _ => None,
+        }
+    }
+
+    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+        Some(ColumnBuilder::Date(builder))
+    }
+
     fn upcast_scalar(scalar: Self::Scalar) -> Scalar {
         Scalar::Date(scalar)
     }

diff --git a/src/query/expression/src/types/decimal.rs b/src/query/expression/src/types/decimal.rs
@@ -89,6 +89,14 @@ impl<Num: Decimal> ValueType for DecimalType<Num> {
         Num::try_downcast_builder(builder)
     }
 
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder> {
+        Num::try_downcast_owned_builder(builder)
+    }
+
+    fn try_upcast_column_builder(_builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+        None
+    }
+
     fn upcast_scalar(scalar: Self::Scalar) -> Scalar {
         Num::upcast_scalar(scalar, Num::default_decimal_size())
     }
@@ -291,6 +299,8 @@ pub trait Decimal:
     fn try_downcast_column(column: &Column) -> Option<(Buffer<Self>, DecimalSize)>;
     fn try_downcast_builder<'a>(builder: &'a mut ColumnBuilder) -> Option<&'a mut Vec<Self>>;
 
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Vec<Self>>;
+
     fn try_downcast_scalar(scalar: &DecimalScalar) -> Option<Self>;
     fn try_downcast_domain(domain: &DecimalDomain) -> Option<SimpleDomain<Self>>;
 
@@ -477,6 +487,13 @@ impl Decimal for i128 {
         }
     }
 
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Vec<Self>> {
+        match builder {
+            ColumnBuilder::Decimal(DecimalColumnBuilder::Decimal128(s, _)) => Some(s),
+            _ => None,
+        }
+    }
+
     fn try_downcast_scalar<'a>(scalar: &DecimalScalar) -> Option<Self> {
         match scalar {
             DecimalScalar::Decimal128(val, _) => Some(*val),
@@ -630,6 +647,13 @@ impl Decimal for i256 {
         }
     }
 
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Vec<Self>> {
+        match builder {
+            ColumnBuilder::Decimal(DecimalColumnBuilder::Decimal256(s, _)) => Some(s),
+            _ => None,
+        }
+    }
+
     fn try_downcast_scalar<'a>(scalar: &DecimalScalar) -> Option<Self> {
         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
@@ -76,6 +76,17 @@ impl ValueType for EmptyArrayType {
         }
     }
 
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder> {
+        match builder {
+            ColumnBuilder::EmptyArray { len } => Some(len),
+            _ => None,
+        }
+    }
+
+    fn try_upcast_column_builder(len: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+        Some(ColumnBuilder::EmptyArray { len })
+    }
+
     fn upcast_scalar(_: Self::Scalar) -> Scalar {
         Scalar::EmptyArray
     }

diff --git a/src/query/expression/src/types/empty_map.rs b/src/query/expression/src/types/empty_map.rs
@@ -76,6 +76,17 @@ impl ValueType for EmptyMapType {
         }
     }
 
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder> {
+        match builder {
+            ColumnBuilder::EmptyMap { len } => Some(len),
+            _ => None,
+        }
+    }
+
+    fn try_upcast_column_builder(len: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+        Some(ColumnBuilder::EmptyMap { len })
+    }
+
     fn upcast_scalar(_: Self::Scalar) -> Scalar {
         Scalar::EmptyMap
     }

diff --git a/src/query/expression/src/types/generic.rs b/src/query/expression/src/types/generic.rs
@@ -67,6 +67,14 @@ impl<const INDEX: usize> ValueType for GenericType<INDEX> {
         Some(builder)
     }
 
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder> {
+        Some(builder)
+    }
+
+    fn try_upcast_column_builder(builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+        Some(builder)
+    }
+
     fn upcast_scalar(scalar: Self::Scalar) -> Scalar {
         scalar
     }

diff --git a/src/query/expression/src/types/map.rs b/src/query/expression/src/types/map.rs
@@ -87,6 +87,14 @@ impl<K: ValueType, V: ValueType> ValueType for KvPair<K, V> {
         None
     }
 
+    fn try_downcast_owned_builder<'a>(_builder: ColumnBuilder) -> Option<Self::ColumnBuilder> {
+        None
+    }
+
+    fn try_upcast_column_builder(_builder: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+        None
+    }
+
     fn upcast_scalar((k, v): Self::Scalar) -> Scalar {
         Scalar::Tuple(vec![K::upcast_scalar(k), V::upcast_scalar(v)])
     }
@@ -351,6 +359,14 @@ impl<K: ValueType, V: ValueType> ValueType for MapType<K, V> {
         <MapInternal<K, V> as ValueType>::try_downcast_builder(builder)
     }
 
+    fn try_downcast_owned_builder<'a>(builder: ColumnBuilder) -> Option<Self::ColumnBuilder> {
+        <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 upcast_scalar(scalar: Self::Scalar) -> Scalar {
         Scalar::Map(KvPair::<K, V>::upcast_column(scalar))
     }

diff --git a/src/query/expression/src/types/null.rs b/src/query/expression/src/types/null.rs
@@ -75,11 +75,22 @@ impl ValueType for NullType {
         builder: &'a mut ColumnBuilder,
     ) -> Option<&'a mut Self::ColumnBuilder> {
         match builder {
-            crate::ColumnBuilder::Null { len } => Some(len),
+            ColumnBuilder::Null { len } => Some(len),
             _ => None,
         }
     }
 
+    fn try_downcast_owned_builder(builder: ColumnBuilder) -> Option<Self::ColumnBuilder> {
+        match builder {
+            ColumnBuilder::Null { len } => Some(len),
+            _ => None,
+        }
+    }
+
+    fn try_upcast_column_builder(len: Self::ColumnBuilder) -> Option<ColumnBuilder> {
+        Some(ColumnBuilder::Null { len })
+    }
+
     fn upcast_scalar(_: Self::Scalar) -> Scalar {
         Scalar::Null
     }