chore(query): update take_mut

src/common/base/Cargo.toml

@@ -48,6 +48,7 @@ pprof = { version = "0.11.1", features = [
     "protobuf",
 ] }
 regex = { workspace = true }
+replace_with = "0.1.7"
 semver = { workspace = true }
 serde = { workspace = true }
 state = "0.5"

src/common/base/src/base/take_mut.rs

@@ -12,29 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::panic;
+use replace_with::replace_with_or_abort_and_return;
 
-use databend_common_exception::Result;
-
-use crate::runtime::catch_unwind;
-
-/// 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 = 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(),
-        }
-    }
+pub fn take_mut<T, U, F: FnOnce(T) -> (U, T)>(dest: &mut T, f: F) -> U {
+    replace_with_or_abort_and_return(dest, f)
 }

src/query/functions/src/aggregates/aggregate_unary.rs

@@ -151,15 +151,21 @@ where
                     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.
-            // See [`take_mut`] for details.
-            _ => 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())
-            }),
+            _ => {
+                // some `ValueType` like `NullableType` need ownership to downcast builder,
+                // so here we using an unsafe way to take the ownership of builder.
+                // See [`take_mut`] for details.
+                if let Some(builder) = R::try_downcast_builder(builder) {
+                    state.merge_result(builder, self.function_data.as_deref())
+                } else {
+                    take_mut(builder, |builder| {
+                        let mut builder = R::try_downcast_owned_builder(builder).unwrap();
+                        let res = state.merge_result(&mut builder, self.function_data.as_deref());
+
+                        (res, R::try_upcast_column_builder(builder).unwrap())
+                    })
+                }
+            }
         }
     }
 }