refactor

Signed-off-by: coldWater <[email protected]>

src/query/expression/src/utils/arrow.rs

@@ -13,17 +13,21 @@
 // limitations under the License.
 
 use std::io::Cursor;
+use std::io::Read;
+use std::io::Seek;
 use std::io::Write;
 
+use databend_common_arrow::arrow;
 use databend_common_arrow::arrow::array::Array;
 use databend_common_arrow::arrow::bitmap::Bitmap;
 use databend_common_arrow::arrow::bitmap::MutableBitmap;
 use databend_common_arrow::arrow::buffer::Buffer;
 use databend_common_arrow::arrow::datatypes::Schema;
 use databend_common_arrow::arrow::io::ipc::read::read_file_metadata;
 use databend_common_arrow::arrow::io::ipc::read::FileReader;
+use databend_common_arrow::arrow::io::ipc::write::Compression;
 use databend_common_arrow::arrow::io::ipc::write::FileWriter;
-use databend_common_arrow::arrow::io::ipc::write::WriteOptions as IpcWriteOptions;
+use databend_common_arrow::arrow::io::ipc::write::WriteOptions;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 
@@ -67,32 +71,31 @@ pub fn buffer_into_mut<T: Clone>(mut buffer: Buffer<T>) -> Vec<T> {
 
 pub fn serialize_column(col: &Column) -> Vec<u8> {
     let mut buffer = Vec::new();
-    serialize_column_in(col, &mut buffer).unwrap();
+    write_column(col, &mut buffer).unwrap();
     buffer
 }
 
-pub fn serialize_column_in(
-    col: &Column,
-    w: &mut impl Write,
-) -> databend_common_arrow::arrow::error::Result<()> {
+pub fn write_column(col: &Column, w: &mut impl Write) -> arrow::error::Result<()> {
     let schema = Schema::from(vec![col.arrow_field()]);
-    let mut writer = FileWriter::new(w, schema, None, IpcWriteOptions::default());
+    let mut writer = FileWriter::new(w, schema, None, WriteOptions {
+        compression: Some(Compression::LZ4),
+    });
     writer.start()?;
-    writer.write(
-        &databend_common_arrow::arrow::chunk::Chunk::new(vec![col.as_arrow()]),
-        None,
-    )?;
+    writer.write(&arrow::chunk::Chunk::new(vec![col.as_arrow()]), None)?;
     writer.finish()
 }
 
 pub fn deserialize_column(bytes: &[u8]) -> Result<Column> {
     let mut cursor = Cursor::new(bytes);
+    read_column(&mut cursor)
+}
 
-    let metadata = read_file_metadata(&mut cursor)?;
+pub fn read_column<R: Read + Seek>(r: &mut R) -> Result<Column> {
+    let metadata = read_file_metadata(r)?;
     let f = metadata.schema.fields[0].clone();
     let data_field = DataField::try_from(&f)?;
 
-    let mut reader = FileReader::new(cursor, metadata, None, None);
+    let mut reader = FileReader::new(r, metadata, None, None);
     let col = reader
         .next()
         .ok_or_else(|| ErrorCode::Internal("expected one arrow array"))??

src/query/service/Cargo.toml

@@ -44,9 +44,11 @@ async-trait = { workspace = true }
 backoff = { version = "0.4.0", features = ["futures", "tokio"] }
 backon = "0.4"
 base64 = "0.21.0"
+buf-list = "1.0.3"
 bumpalo = { workspace = true }
 byte-unit = "4.0.19"
 byteorder = { workspace = true }
+bytes = { workspace = true }
 chrono = { workspace = true }
 chrono-tz = { workspace = true }
 config = { version = "0.13.4", features = [] }

src/query/service/src/pipelines/builders/builder_window.rs

@@ -26,6 +26,8 @@ use databend_common_pipeline_core::processors::ProcessorPtr;
 use databend_common_sql::executor::physical_plans::Window;
 use databend_common_sql::executor::physical_plans::WindowPartition;
 use databend_storages_common_cache::TempDirManager;
+use opendal::services::Fs;
+use opendal::Operator;
 
 use crate::pipelines::processors::transforms::FrameBound;
 use crate::pipelines::processors::transforms::TransformWindowPartitionCollect;
@@ -34,6 +36,7 @@ use crate::pipelines::processors::transforms::WindowPartitionExchange;
 use crate::pipelines::processors::transforms::WindowSpillSettings;
 use crate::pipelines::processors::TransformWindow;
 use crate::pipelines::PipelineBuilder;
+use crate::spillers::SpillerDiskConfig;
 
 impl PipelineBuilder {
     pub(crate) fn build_window(&mut self, window: &Window) -> Result<()> {
@@ -173,7 +176,23 @@ impl PipelineBuilder {
 
         let disk_bytes_limit = settings.get_window_partition_spilling_to_disk_bytes_limit()?;
         let temp_dir_manager = TempDirManager::instance();
-        let disk_spill = temp_dir_manager.get_disk_spill_dir(disk_bytes_limit, &self.ctx.get_id());
+
+        let enable_dio = settings.get_enable_dio()?;
+        let disk_spill =
+            match temp_dir_manager.get_disk_spill_dir(disk_bytes_limit, &self.ctx.get_id()) {
+                Some(temp_dir) if !enable_dio => {
+                    let builder = Fs::default().root(temp_dir.path().to_str().unwrap());
+                    Some(SpillerDiskConfig {
+                        temp_dir,
+                        local_operator: Some(Operator::new(builder)?.finish()),
+                    })
+                }
+                Some(temp_dir) => Some(SpillerDiskConfig {
+                    temp_dir,
+                    local_operator: None,
+                }),
+                None => None,
+            };
 
         let window_spill_settings = WindowSpillSettings::new(&settings, num_processors)?;
         let have_order_col = window_partition.after_exchange.unwrap_or(false);

src/query/service/src/pipelines/processors/transforms/window/partition/transform_window_partition_collect.rs

@@ -34,14 +34,14 @@ use databend_common_pipeline_transforms::processors::sort_merge;
 use databend_common_settings::Settings;
 use databend_common_storage::DataOperator;
 use databend_common_storages_fuse::TableContext;
-use databend_storages_common_cache::TempDir;
 
 use super::WindowPartitionBuffer;
 use super::WindowPartitionMeta;
 use super::WindowSpillSettings;
 use crate::sessions::QueryContext;
 use crate::spillers::Spiller;
 use crate::spillers::SpillerConfig;
+use crate::spillers::SpillerDiskConfig;
 use crate::spillers::SpillerType;
 
 #[derive(Debug, Clone, Copy)]
@@ -99,7 +99,7 @@ impl TransformWindowPartitionCollect {
         num_processors: usize,
         num_partitions: usize,
         spill_settings: WindowSpillSettings,
-        disk_spill: Option<Arc<TempDir>>,
+        disk_spill: Option<SpillerDiskConfig>,
         sort_desc: Vec<SortColumnDescription>,
         schema: DataSchemaRef,
         have_order_col: bool,

src/query/service/src/spillers/mod.rs

@@ -17,8 +17,4 @@ mod spiller;
 
 pub use partition_buffer::PartitionBuffer;
 pub use partition_buffer::PartitionBufferFetchOption;
-pub use spiller::Location;
-pub use spiller::SpilledData;
-pub use spiller::Spiller;
-pub use spiller::SpillerConfig;
-pub use spiller::SpillerType;
+pub use spiller::*;