split memory manager to module auron-memmgr

Cargo.toml

@@ -24,6 +24,7 @@ members = [
     "native-engine/auron",
     "native-engine/auron-jni-bridge",
     "native-engine/auron-serde",
+    "native-engine/auron-memmgr",
 ]
 
 [profile.release]
@@ -48,6 +49,7 @@ overflow-checks = false
 auron = { path = "./native-engine/auron" }
 auron-jni-bridge = { path = "./native-engine/auron-jni-bridge" }
 auron-serde = { path = "./native-engine/auron-serde" }
+auron-memmgr = { path = "./native-engine/auron-memmgr" }
 datafusion-ext-commons = { path = "./native-engine/datafusion-ext-commons" }
 datafusion-ext-exprs = { path = "./native-engine/datafusion-ext-exprs" }
 datafusion-ext-functions = { path = "./native-engine/datafusion-ext-functions" }

native-engine/auron-memmgr/Cargo.toml

@@ -0,0 +1,37 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You 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.
+#
+
+[package]
+name = "auron-memmgr"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+auron-jni-bridge = { workspace = true }
+datafusion = { workspace = true }
+datafusion-ext-commons = { workspace = true }
+
+async-trait = { workspace = true }
+bytesize = { workspace = true }
+jni = { workspace = true }
+log = { workspace = true }
+once_cell = { workspace = true }
+tempfile = { workspace = true }
+parking_lot = { workspace = true }
+
+[target.'cfg(target_os = "linux")'.dependencies]
+procfs = { workspace = true }

native-engine/datafusion-ext-plans/src/memmgr/mod.rs → native-engine/auron-memmgr/src/lib.rs

@@ -13,6 +13,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#![feature(get_mut_unchecked)]
+
 pub mod metrics;
 pub mod spill;
 

native-engine/datafusion-ext-plans/src/memmgr/spill.rs → native-engine/auron-memmgr/src/spill.rs

@@ -27,14 +27,12 @@ use auron_jni_bridge::{
     jni_get_string, jni_new_direct_byte_buffer, jni_new_global_ref,
 };
 use datafusion::{common::Result, parquet::file::reader::Length, physical_plan::metrics::Time};
+use datafusion_ext_commons::io::ipc_compression::{IoCompressionReader, IoCompressionWriter};
 use jni::{objects::GlobalRef, sys::jlong};
 use log::warn;
 use once_cell::sync::OnceCell;
 
-use crate::{
-    common::ipc_compression::{IoCompressionReader, IoCompressionWriter},
-    memmgr::metrics::SpillMetrics,
-};
+use crate::metrics::SpillMetrics;
 
 pub type SpillCompressedReader<'a> = IoCompressionReader<BufReader<Box<dyn Read + Send + 'a>>>;
 pub type SpillCompressedWriter<'a> = IoCompressionWriter<BufWriter<Box<dyn Write + Send + 'a>>>;

native-engine/auron-serde/Cargo.toml

@@ -26,17 +26,15 @@ default = ["prost/no-recursion-limit"]
 [dependencies]
 arrow = { workspace = true }
 datafusion = { workspace = true }
-datafusion-ext-commons = { workspace = true }
 datafusion-ext-exprs = { workspace = true }
 datafusion-ext-functions = { workspace = true }
 datafusion-ext-plans = { workspace = true }
 datafusion-spark = { workspace = true }
 
 base64 = { workspace = true }
-log = { workspace = true }
 object_store = { workspace = true }
-prost = { workspace = true }
 parking_lot = { workspace = true }
+prost = { workspace = true }
 
 [build-dependencies]
 tonic-build = { workspace = true }

native-engine/auron/Cargo.toml

@@ -33,12 +33,12 @@ http-service = []
 [dependencies]
 arrow = { workspace = true }
 auron-jni-bridge = { workspace = true }
+auron-memmgr = { workspace = true }
 auron-serde = { workspace = true }
 datafusion = { workspace = true }
 datafusion-ext-commons = { workspace = true }
 datafusion-ext-plans = { workspace = true }
 
-bytesize = { workspace = true }
 futures = { workspace = true }
 jni = { workspace = true }
 log = { workspace = true }

native-engine/auron/src/exec.rs

@@ -18,6 +18,7 @@ use auron_jni_bridge::{
     jni_bridge::JavaClasses,
     *,
 };
+use auron_memmgr::MemManager;
 use datafusion::{
     common::Result,
     error::DataFusionError,
@@ -28,7 +29,6 @@ use datafusion::{
     },
     prelude::{SessionConfig, SessionContext},
 };
-use datafusion_ext_plans::memmgr::MemManager;
 use jni::{
     JNIEnv,
     objects::{JClass, JObject, JString},

native-engine/datafusion-ext-commons/Cargo.toml

@@ -29,24 +29,21 @@ arrow = { workspace = true }
 arrow-schema = { workspace = true }
 auron-jni-bridge = { workspace = true }
 datafusion = { workspace = true }
-
-async-trait = { workspace = true }
 bigdecimal = { workspace = true }
 byteorder = { workspace = true }
-bytes = { workspace = true }
 chrono = { workspace = true }
-futures = { workspace = true }
 itertools = { workspace = true }
 jni = { workspace = true }
 log = { workspace = true }
+lz4_flex = { workspace = true }
 num = { workspace = true }
 once_cell = { workspace = true }
 paste = { workspace = true }
 smallvec = { workspace = true }
-tempfile = { workspace = true }
-transpose = { workspace = true }
 tokio = { workspace = true }
+transpose = { workspace = true }
 unchecked-index = { workspace = true }
+zstd = { workspace = true }
 
 [dev-dependencies]
 rand = { workspace = true }

native-engine/datafusion-ext-plans/src/common/ipc_compression.rs → native-engine/datafusion-ext-commons/src/io/ipc_compression.rs

@@ -25,11 +25,12 @@ use auron_jni_bridge::{
 };
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 use datafusion::common::Result;
-use datafusion_ext_commons::{
+use once_cell::sync::OnceCell;
+
+use crate::{
     df_execution_err,
     io::{read_one_batch, write_one_batch},
 };
-use once_cell::sync::OnceCell;
 
 pub struct IpcCompressionWriter<W: Write> {
     output: W,

native-engine/datafusion-ext-commons/src/io/mod.rs

@@ -27,6 +27,7 @@ pub use scalar_serde::{read_scalar, write_scalar};
 use crate::{UninitializedInit, arrow::cast::cast};
 
 mod batch_serde;
+pub mod ipc_compression;
 mod scalar_serde;
 
 pub fn write_one_batch(num_rows: usize, cols: &[ArrayRef], mut output: impl Write) -> Result<()> {

native-engine/datafusion-ext-exprs/Cargo.toml

@@ -23,14 +23,12 @@ resolver = "1"
 
 [dependencies]
 arrow = { workspace = true }
-async-trait = "0.1.89"
 auron-jni-bridge = { workspace = true }
 datafusion = { workspace = true }
 datafusion-ext-commons = { workspace = true }
 
 itertools = { workspace = true }
 jni = { workspace = true }
 log = { workspace = true }
-num = { workspace = true }
 once_cell = { workspace = true }
 parking_lot = { workspace = true }

native-engine/datafusion-ext-functions/Cargo.toml

@@ -23,7 +23,6 @@ resolver = "1"
 
 [dependencies]
 arrow = { workspace = true }
-async-trait = "0.1.89"
 auron-jni-bridge = { workspace = true }
 datafusion = { workspace = true }
 datafusion-ext-commons = { workspace = true }
@@ -33,4 +32,4 @@ log = { workspace = true }
 num = { workspace = true }
 paste = { workspace = true }
 serde_json = { workspace = true }
-sonic-rs = { workspace = true }
+sonic-rs = { workspace = true }

native-engine/datafusion-ext-plans/Cargo.toml

@@ -28,6 +28,7 @@ default = ["tokio/rt-multi-thread"]
 arrow = { workspace = true }
 arrow-schema = { workspace = true }
 auron-jni-bridge = { workspace = true }
+auron-memmgr = { workspace = true }
 datafusion = { workspace = true }
 datafusion-datasource = { workspace = true }
 datafusion-datasource-parquet = { workspace = true }
@@ -50,18 +51,15 @@ hashbrown = { workspace = true }
 itertools = { workspace = true }
 jni = { workspace = true }
 log = { workspace = true }
-lz4_flex = { workspace = true }
 num = { workspace = true }
 object_store = { workspace = true }
 once_cell = { workspace = true }
 panic-message = { workspace = true }
 parking_lot = { workspace = true }
 paste = { workspace = true }
 smallvec = { workspace = true }
-tempfile = { workspace = true }
 tokio = { workspace = true }
 unchecked-index = { workspace = true }
-zstd = { workspace = true }
 
 [target.'cfg(target_os = "linux")'.dependencies]
 procfs = { workspace = true }

native-engine/datafusion-ext-plans/src/agg/acc.rs

@@ -23,6 +23,7 @@ use arrow::{
     array::*,
     datatypes::{DataType, *},
 };
+use auron_memmgr::spill::{SpillCompressedReader, SpillCompressedWriter};
 use bitvec::{bitvec, vec::BitVec};
 use byteorder::{ReadBytesExt, WriteBytesExt};
 use datafusion::common::{Result, ScalarValue, utils::proxy::VecAllocExt};
@@ -33,11 +34,7 @@ use datafusion_ext_commons::{
 };
 use smallvec::SmallVec;
 
-use crate::{
-    agg::agg::IdxSelection,
-    idx_for, idx_with_iter,
-    memmgr::spill::{SpillCompressedReader, SpillCompressedWriter},
-};
+use crate::{agg::agg::IdxSelection, idx_for, idx_with_iter};
 
 pub trait AccColumn: Send {
     fn as_any(&self) -> &dyn Any;