refactor(query): split Spiller to provide more scalability (#18691)

* fmt

* spiller

* BroadcastChannel

* move

Author: forsaken628

Cargo.lock

@@ -15,6 +15,7 @@ databend-common-column = { workspace = true }
 databend-common-exception = { workspace = true }
 databend-common-expression = { workspace = true }
 databend-common-pipeline-core = { workspace = true }
+databend-storages-common-cache = { workspace = true }
 jsonb = { workspace = true }
 log = { workspace = true }
 match-template = { workspace = true }

src/query/pipeline/transforms/Cargo.toml

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 mod memory_settings;
+pub mod traits;
 mod transforms;
 
 pub use memory_settings::MemorySettings;

src/query/pipeline/transforms/src/processors/mod.rs

@@ -0,0 +1,17 @@
+// 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.
+
+mod spill;
+
+pub use spill::*;

src/query/pipeline/transforms/src/processors/traits/mod.rs

@@ -0,0 +1,32 @@
+// 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 databend_common_exception::Result;
+use databend_common_expression::DataBlock;
+use databend_storages_common_cache::TempPath;
+
+#[derive(Debug, Clone, Hash, PartialEq, Eq)]
+pub enum Location {
+    Remote(String),
+    Local(TempPath),
+}
+
+#[async_trait::async_trait]
+pub trait DataBlockSpill: Clone + Send + Sync + 'static {
+    async fn spill(&self, data_block: DataBlock) -> Result<Location> {
+        self.merge_and_spill(vec![data_block]).await
+    }
+    async fn merge_and_spill(&self, data_block: Vec<DataBlock>) -> Result<Location>;
+    async fn restore(&self, location: &Location) -> Result<DataBlock>;
+}

src/query/pipeline/transforms/src/processors/traits/spill.rs

@@ -49,8 +49,11 @@ use crate::physical_plans::WindowPartition;
 use crate::physical_plans::WindowPartitionTopN;
 use crate::physical_plans::WindowPartitionTopNFunc;
 use crate::pipelines::builders::SortPipelineBuilder;
-use crate::pipelines::processors::transforms::TransformSortBuilder;
 use crate::pipelines::PipelineBuilder;
+use crate::spillers::SortSpiller;
+
+type TransformSortBuilder =
+    crate::pipelines::processors::transforms::TransformSortBuilder<SortSpiller>;
 
 #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
 pub struct Sort {

src/query/service/src/physical_plans/physical_sort.rs

@@ -35,13 +35,15 @@ use databend_common_storages_fuse::TableContext;
 use databend_storages_common_cache::TempDirManager;
 
 use crate::pipelines::memory_settings::MemorySettingsExt;
-use crate::pipelines::processors::transforms::TransformSortBuilder;
 use crate::sessions::QueryContext;
-use crate::spillers::Spiller;
+use crate::spillers::SortSpiller;
 use crate::spillers::SpillerConfig;
 use crate::spillers::SpillerDiskConfig;
 use crate::spillers::SpillerType;
 
+type TransformSortBuilder =
+    crate::pipelines::processors::transforms::TransformSortBuilder<SortSpiller>;
+
 pub struct SortPipelineBuilder {
     ctx: Arc<QueryContext>,
     output_schema: DataSchemaRef,
@@ -145,7 +147,7 @@ impl SortPipelineBuilder {
                 use_parquet: settings.get_spilling_file_format()?.is_parquet(),
             };
             let op = DataOperator::instance().spill_operator();
-            Arc::new(Spiller::create(self.ctx.clone(), op, config)?)
+            SortSpiller::new(self.ctx.clone(), op, config)?
         };
 
         pipeline.add_transform(|input, output| {
@@ -231,7 +233,7 @@ impl SortPipelineBuilder {
                 use_parquet: settings.get_spilling_file_format()?.is_parquet(),
             };
             let op = DataOperator::instance().spill_operator();
-            Arc::new(Spiller::create(self.ctx.clone(), op, config)?)
+            SortSpiller::new(self.ctx.clone(), op, config)?
         };
 
         let memory_settings = MemorySettings::from_sort_settings(&self.ctx)?;

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

@@ -162,7 +162,6 @@ impl PipelinePullingExecutor {
         })
     }
 
-    #[fastrace::trace]
     pub fn start(&mut self) {
         let _guard = ThreadTracker::tracking(self.tracking_payload.clone());
 

src/query/service/src/pipelines/executor/pipeline_pulling_executor.rs

@@ -31,6 +31,7 @@ use databend_common_pipeline_core::processors::Processor;
 use databend_common_pipeline_transforms::processors::sort::algorithm::SortAlgorithm;
 use databend_common_pipeline_transforms::sort::RowConverter;
 use databend_common_pipeline_transforms::sort::Rows;
+use databend_common_pipeline_transforms::traits::DataBlockSpill;
 use databend_common_pipeline_transforms::MemorySettings;
 use databend_common_pipeline_transforms::MergeSort;
 use databend_common_pipeline_transforms::SortSpillParams;
@@ -42,7 +43,6 @@ use super::sort_spill::OutputData;
 use super::sort_spill::SortSpill;
 use super::Base;
 use super::RowsStat;
-use crate::spillers::Spiller;
 
 #[derive(Debug)]
 enum State {
@@ -54,14 +54,14 @@ enum State {
     Finish,
 }
 
-enum Inner<A: SortAlgorithm> {
+enum Inner<A: SortAlgorithm, S: DataBlockSpill> {
     Collect(Vec<DataBlock>),
     Limit(TransformSortMergeLimit<A::Rows>),
     Memory(MemoryMerger<A>),
-    Spill(Vec<DataBlock>, SortSpill<A>),
+    Spill(Vec<DataBlock>, SortSpill<A, S>),
 }
 
-pub struct TransformSort<A: SortAlgorithm, C> {
+pub struct TransformSort<A: SortAlgorithm, C, S: DataBlockSpill> {
     name: &'static str,
     input: Arc<InputPort>,
     output: Arc<OutputPort>,
@@ -79,19 +79,20 @@ pub struct TransformSort<A: SortAlgorithm, C> {
     /// so we don't need to generate the order column again.
     order_col_generated: bool,
 
-    base: Base,
-    inner: Inner<A>,
+    base: Base<S>,
+    inner: Inner<A, S>,
 
     aborting: AtomicBool,
 
     max_block_size: usize,
     memory_settings: MemorySettings,
 }
 
-impl<A, C> TransformSort<A, C>
+impl<A, C, S> TransformSort<A, C, S>
 where
     A: SortAlgorithm,
     C: RowConverter<A::Rows>,
+    S: DataBlockSpill,
 {
     pub(super) fn new(
         input: Arc<InputPort>,
@@ -100,7 +101,7 @@ where
         sort_desc: Arc<[SortColumnDescription]>,
         max_block_size: usize,
         limit: Option<(usize, bool)>,
-        spiller: Arc<Spiller>,
+        spiller: S,
         output_order_col: bool,
         order_col_generated: bool,
         memory_settings: MemorySettings,
@@ -303,11 +304,12 @@ where
 }
 
 #[async_trait::async_trait]
-impl<A, C> Processor for TransformSort<A, C>
+impl<A, C, S> Processor for TransformSort<A, C, S>
 where
     A: SortAlgorithm + 'static,
     A::Rows: 'static,
     C: RowConverter<A::Rows> + Send + 'static,
+    S: DataBlockSpill,
 {
     fn name(&self) -> String {
         self.name.to_string()

src/query/service/src/pipelines/processors/transforms/sort/merge_sort.rs

@@ -12,8 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::sync::Arc;
-
 use bounds::Bounds;
 use databend_common_expression::local_block_meta_serde;
 use databend_common_expression::BlockMetaInfo;
@@ -24,8 +22,6 @@ use databend_common_pipeline_transforms::SortSpillParams;
 use enum_as_inner::EnumAsInner;
 use sort_spill::SpillableBlock;
 
-use crate::spillers::Spiller;
-
 mod bounds;
 mod merge_sort;
 mod sort_broadcast;
@@ -36,8 +32,6 @@ mod sort_merge_stream;
 mod sort_restore;
 mod sort_route;
 mod sort_spill;
-#[cfg(test)]
-mod test_memory;
 
 pub use merge_sort::*;
 pub use sort_broadcast::*;
@@ -49,9 +43,9 @@ pub use sort_restore::*;
 pub use sort_route::*;
 
 #[derive(Clone)]
-struct Base {
+struct Base<S: Send + Clone> {
     schema: DataSchemaRef,
-    spiller: Arc<Spiller>,
+    spiller: S,
     sort_row_offset: usize,
     limit: Option<usize>,
 }