add compact strategy

Author: zhyass

src/query/service/src/pipelines/processors/transforms/mod.rs

@@ -17,6 +17,7 @@ pub mod aggregator;
 mod broadcast;
 mod hash_join;
 pub(crate) mod range_join;
+mod recluster;
 mod runtime_pool;
 mod transform_add_computed_columns;
 mod transform_add_const_columns;
@@ -45,6 +46,7 @@ mod window;
 pub use broadcast::BroadcastSinkProcessor;
 pub use broadcast::BroadcastSourceProcessor;
 pub use hash_join::*;
+pub use recluster::*;
 pub use transform_add_computed_columns::TransformAddComputedColumns;
 pub use transform_add_const_columns::TransformAddConstColumns;
 pub use transform_add_internal_columns::TransformAddInternalColumns;

src/query/service/src/pipelines/processors/transforms/recluster/compact_strategy.rs

@@ -0,0 +1,78 @@
+// 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 crate::pipelines::processors::transforms::DataProcessorStrategy;
+
+pub struct CompactStrategy {
+    max_bytes_per_block: usize,
+    max_rows_per_block: usize,
+}
+
+impl CompactStrategy {
+    pub fn new(max_rows_per_block: usize, max_bytes_per_block: usize) -> Self {
+        Self {
+            max_bytes_per_block,
+            max_rows_per_block,
+        }
+    }
+
+    fn concat_blocks(blocks: Vec<DataBlock>) -> Result<DataBlock> {
+        DataBlock::concat(&blocks)
+    }
+
+    fn check_large_enough(&self, rows: usize, bytes: usize) -> bool {
+        rows >= self.max_rows_per_block || bytes >= self.max_bytes_per_block
+    }
+}
+
+impl DataProcessorStrategy for CompactStrategy {
+    const NAME: &'static str = "Compact";
+
+    fn process_data_blocks(&self, data_blocks: Vec<DataBlock>) -> Result<Vec<DataBlock>> {
+        let blocks_num = data_blocks.len();
+        if blocks_num < 2 {
+            return Ok(data_blocks);
+        }
+
+        let mut accumulated_rows = 0;
+        let mut accumulated_bytes = 0;
+        let mut pending_blocks = Vec::with_capacity(blocks_num);
+        let mut staged_blocks = Vec::with_capacity(blocks_num);
+        let mut result = Vec::with_capacity(blocks_num);
+        for block in data_blocks {
+            accumulated_rows += block.num_rows();
+            accumulated_bytes += block.estimate_block_size();
+            pending_blocks.push(block);
+            if !self.check_large_enough(accumulated_rows, accumulated_bytes) {
+                continue;
+            }
+            if !staged_blocks.is_empty() {
+                result.push(Self::concat_blocks(std::mem::take(&mut staged_blocks))?);
+            }
+            std::mem::swap(&mut staged_blocks, &mut pending_blocks);
+            accumulated_rows = 0;
+            accumulated_bytes = 0;
+        }
+
+        staged_blocks.append(&mut pending_blocks);
+        if !staged_blocks.is_empty() {
+            result.push(Self::concat_blocks(std::mem::take(&mut staged_blocks))?);
+        }
+
+        Ok(result)
+    }
+}

src/query/service/src/pipelines/processors/transforms/window/partition/hilbert_partition_exchange.rs renamed to src/query/service/src/pipelines/processors/transforms/recluster/hilbert_partition_exchange.rs

@@ -12,10 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Some variables and functions are named and designed with reference to ClickHouse.
-// - https://github.com/ClickHouse/ClickHouse/blob/master/src/Processors/Transforms/WindowTransform.h
-// - https://github.com/ClickHouse/ClickHouse/blob/master/src/Processors/Transforms/WindowTransform.cpp
-
 use std::sync::Arc;
 
 use databend_common_exception::Result;

src/query/service/src/pipelines/processors/transforms/recluster/mod.rs

@@ -0,0 +1,21 @@
+// 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 compact_strategy;
+mod hilbert_partition_exchange;
+mod transform_hilbert_collect;
+
+pub use compact_strategy::CompactStrategy;
+pub use hilbert_partition_exchange::HilbertPartitionExchange;
+pub use transform_hilbert_collect::TransformHilbertCollect;

src/query/service/src/pipelines/processors/transforms/window/partition/transform_hilbert_collect.rs renamed to src/query/service/src/pipelines/processors/transforms/recluster/transform_hilbert_collect.rs

@@ -27,8 +27,10 @@ use databend_common_pipeline_transforms::MemorySettings;
 use databend_common_settings::Settings;
 use databend_common_storage::DataOperator;
 
-use super::WindowPartitionBuffer;
-use super::WindowPartitionMeta;
+use crate::pipelines::processors::transforms::CompactStrategy;
+use crate::pipelines::processors::transforms::DataProcessorStrategy;
+use crate::pipelines::processors::transforms::WindowPartitionBuffer;
+use crate::pipelines::processors::transforms::WindowPartitionMeta;
 use crate::sessions::QueryContext;
 use crate::spillers::Spiller;
 use crate::spillers::SpillerConfig;
@@ -40,7 +42,7 @@ enum State {
     Flush,
     Spill,
     Restore,
-    Concat(Vec<DataBlock>),
+    Compact(Vec<DataBlock>),
 }
 
 pub struct TransformHilbertCollect {
@@ -56,6 +58,7 @@ pub struct TransformHilbertCollect {
     // The buffer is used to control the memory usage of the window operator.
     buffer: WindowPartitionBuffer,
 
+    compact_strategy: CompactStrategy,
     max_block_size: usize,
     // Event variables.
     state: State,
@@ -100,8 +103,6 @@ impl TransformHilbertCollect {
         let spiller = Spiller::create(ctx, operator, spill_config)?;
 
         // Create the window partition buffer.
-        let max_block_rows =
-            max_block_rows.min(settings.get_window_partition_sort_block_size()? as usize);
         let buffer =
             WindowPartitionBuffer::new(spiller, partitions.len(), max_block_rows, memory_settings)?;
 
@@ -113,6 +114,7 @@ impl TransformHilbertCollect {
             immediate_output_blocks: vec![],
             partition_sizes: vec![0; num_partitions],
             max_block_size,
+            compact_strategy: CompactStrategy::new(max_block_rows, max_block_size),
             output_data_blocks: VecDeque::new(),
             state: State::Collect,
         })
@@ -130,7 +132,7 @@ impl Processor for TransformHilbertCollect {
     }
 
     fn event(&mut self) -> Result<Event> {
-        if matches!(self.state, State::Concat(_)) {
+        if matches!(self.state, State::Compact(_)) {
             return Ok(Event::Sync);
         }
 
@@ -192,9 +194,9 @@ impl Processor for TransformHilbertCollect {
 
     fn process(&mut self) -> Result<()> {
         match std::mem::replace(&mut self.state, State::Collect) {
-            State::Concat(blocks) => {
-                let output = DataBlock::concat(&blocks)?;
-                self.output_data_blocks.push_back(output);
+            State::Compact(blocks) => {
+                let output = self.compact_strategy.process_data_blocks(blocks)?;
+                self.output_data_blocks.extend(output);
             }
             _ => unreachable!(),
         }
@@ -212,12 +214,12 @@ impl Processor for TransformHilbertCollect {
                     let mut restored_data_blocks =
                         self.buffer.restore_by_id(partition_id, true).await?;
                     restored_data_blocks.push(data_block);
-                    self.state = State::Concat(restored_data_blocks);
+                    self.state = State::Compact(restored_data_blocks);
                 }
             }
             State::Restore => {
                 let restored_data_blocks = self.buffer.restore().await?;
-                self.output_data_blocks.extend(restored_data_blocks);
+                self.state = State::Compact(restored_data_blocks);
             }
             _ => unreachable!(),
         }

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

@@ -24,66 +24,6 @@ pub trait DataProcessorStrategy: Send + Sync + 'static {
     fn process_data_blocks(&self, data_blocks: Vec<DataBlock>) -> Result<Vec<DataBlock>>;
 }
 
-pub struct CompactStrategy {
-    max_bytes_per_block: usize,
-    max_rows_per_block: usize,
-}
-
-impl CompactStrategy {
-    pub fn new(max_rows_per_block: usize, max_bytes_per_block: usize) -> Self {
-        Self {
-            max_bytes_per_block,
-            max_rows_per_block,
-        }
-    }
-
-    fn concat_blocks(blocks: Vec<DataBlock>) -> Result<DataBlock> {
-        DataBlock::concat(&blocks)
-    }
-
-    fn check_large_enough(&self, rows: usize, bytes: usize) -> bool {
-        rows >= self.max_rows_per_block || bytes >= self.max_bytes_per_block
-    }
-}
-
-impl DataProcessorStrategy for CompactStrategy {
-    const NAME: &'static str = "Compact";
-
-    fn process_data_blocks(&self, data_blocks: Vec<DataBlock>) -> Result<Vec<DataBlock>> {
-        let blocks_num = data_blocks.len();
-        if blocks_num < 2 {
-            return Ok(data_blocks);
-        }
-
-        let mut accumulated_rows = 0;
-        let mut accumulated_bytes = 0;
-        let mut pending_blocks = Vec::with_capacity(blocks_num);
-        let mut staged_blocks = Vec::with_capacity(blocks_num);
-        let mut result = Vec::with_capacity(blocks_num);
-        for block in data_blocks {
-            accumulated_rows += block.num_rows();
-            accumulated_bytes += block.estimate_block_size();
-            pending_blocks.push(block);
-            if !self.check_large_enough(accumulated_rows, accumulated_bytes) {
-                continue;
-            }
-            if !staged_blocks.is_empty() {
-                result.push(Self::concat_blocks(std::mem::take(&mut staged_blocks))?);
-            }
-            std::mem::swap(&mut staged_blocks, &mut pending_blocks);
-            accumulated_rows = 0;
-            accumulated_bytes = 0;
-        }
-
-        staged_blocks.append(&mut pending_blocks);
-        if !staged_blocks.is_empty() {
-            result.push(Self::concat_blocks(std::mem::take(&mut staged_blocks))?);
-        }
-
-        Ok(result)
-    }
-}
-
 pub struct SortStrategy {
     sort_desc: Vec<SortColumnDescription>,
     schema: DataSchemaRef,

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

@@ -13,17 +13,13 @@
 // limitations under the License.
 
 mod data_processor_strategy;
-mod hilbert_partition_exchange;
-mod transform_hilbert_collect;
 mod transform_window_partition_collect;
 mod window_partition_buffer;
 mod window_partition_exchange;
 mod window_partition_meta;
 mod window_partition_partial_top_n_exchange;
 
 pub use data_processor_strategy::*;
-pub use hilbert_partition_exchange::*;
-pub use transform_hilbert_collect::*;
 pub use transform_window_partition_collect::*;
 pub use window_partition_buffer::*;
 pub use window_partition_exchange::*;