feat(query): Support use parquet format when spilling

src/query/ast/src/parser/parser.rs

@@ -145,6 +145,7 @@ pub fn run_parser<O>(
 }
 
 /// Check that the statement can be displayed and reparsed without loss
+#[allow(dead_code)]
 fn assert_reparse(sql: &str, stmt: StatementWithFormat) {
     let stmt = reset_ast(stmt);
     let new_sql = stmt.to_string();
@@ -162,6 +163,7 @@ fn assert_reparse(sql: &str, stmt: StatementWithFormat) {
     assert_eq!(stmt, new_stmt, "\nleft:\n{}\nright:\n{}", sql, new_sql);
 }
 
+#[allow(dead_code)]
 fn reset_ast(mut stmt: StatementWithFormat) -> StatementWithFormat {
     #[derive(VisitorMut)]
     #[visitor(Range(enter), Literal(enter), ExplainKind(enter), SelectTarget(enter))]

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

@@ -283,12 +283,13 @@ impl SortPipelineBuilder {
         if may_spill {
             let schema = add_order_field(sort_merge_output_schema.clone(), &self.sort_desc);
             let config = SpillerConfig {
+                spiller_type: SpillerType::OrderBy,
                 location_prefix: query_spill_prefix(
                     self.ctx.get_tenant().tenant_name(),
                     &self.ctx.get_id(),
                 ),
                 disk_spill: None,
-                spiller_type: SpillerType::OrderBy,
+                use_parquet: settings.get_spilling_file_format()?.is_parquet(),
             };
             pipeline.add_transform(|input, output| {
                 let op = DataOperator::instance().operator();

src/query/service/src/pipelines/processors/transforms/hash_join/hash_join_spiller.rs

@@ -68,9 +68,10 @@ impl HashJoinSpiller {
             SpillerType::HashJoinProbe
         };
         let spill_config = SpillerConfig {
+            spiller_type,
             location_prefix: query_spill_prefix(ctx.get_tenant().tenant_name(), &ctx.get_id()),
             disk_spill: None,
-            spiller_type,
+            use_parquet: ctx.get_settings().get_spilling_file_format()?.is_parquet(),
         };
         let operator = DataOperator::instance().operator();
         let spiller = Spiller::create(ctx.clone(), operator, spill_config)?;
@@ -145,9 +146,8 @@ impl HashJoinSpiller {
                     .partition_buffer
                     .fetch_data_blocks(partition_id, &fetch_option)?
                 {
-                    let data_block = DataBlock::concat(&data_blocks)?;
                     self.spiller
-                        .spill_with_partition(partition_id, data_block)
+                        .spill_with_partition(partition_id, data_blocks)
                         .await?;
                 }
             }

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

@@ -270,7 +270,7 @@ where R: Rows + Sync + Send + 'static
     async fn spill(&mut self, block: DataBlock) -> Result<()> {
         debug_assert!(self.num_merge >= 2 && self.batch_rows > 0);
 
-        let location = self.spiller.spill(block).await?;
+        let location = self.spiller.spill(vec![block]).await?;
 
         self.unmerged_blocks.push_back(vec![location].into());
         Ok(())
@@ -347,7 +347,7 @@ where R: Rows + Sync + Send + 'static
 
         let mut spilled = VecDeque::new();
         while let Some(block) = merger.async_next_block().await? {
-            let location = self.spiller.spill(block).await?;
+            let location = self.spiller.spill(vec![block]).await?;
 
             spilled.push_back(location);
         }
@@ -456,6 +456,7 @@ mod tests {
     use std::sync::Arc;
 
     use databend_common_base::base::tokio;
+    use databend_common_catalog::table_context::TableContext;
     use databend_common_exception::Result;
     use databend_common_expression::block_debug::pretty_format_blocks;
     use databend_common_expression::types::DataType;
@@ -487,9 +488,10 @@ mod tests {
     ) -> Result<TransformSortSpill<SimpleRowsAsc<Int32Type>>> {
         let op = DataOperator::instance().operator();
         let spill_config = SpillerConfig {
+            spiller_type: SpillerType::OrderBy,
             location_prefix: "_spill_test".to_string(),
             disk_spill: None,
-            spiller_type: SpillerType::OrderBy,
+            use_parquet: ctx.get_settings().get_spilling_file_format()?.is_parquet(),
         };
 
         let spiller = Spiller::create(ctx.clone(), op, spill_config)?;

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

@@ -116,9 +116,10 @@ impl TransformWindowPartitionCollect {
         }
 
         let spill_config = SpillerConfig {
+            spiller_type: SpillerType::Window,
             location_prefix: query_spill_prefix(ctx.get_tenant().tenant_name(), &ctx.get_id()),
             disk_spill,
-            spiller_type: SpillerType::Window,
+            use_parquet: settings.get_spilling_file_format()?.is_parquet(),
         };
 
         // Create an inner `Spiller` to spill data.

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

@@ -17,9 +17,9 @@ use databend_common_exception::Result;
 use databend_common_expression::DataBlock;
 use databend_common_settings::Settings;
 
+use crate::spillers::MergedPartition;
 use crate::spillers::PartitionBuffer;
 use crate::spillers::PartitionBufferFetchOption;
-use crate::spillers::SpilledData;
 use crate::spillers::Spiller;
 
 /// The `WindowPartitionBuffer` is used to control memory usage of Window operator.
@@ -33,7 +33,7 @@ pub struct WindowPartitionBuffer {
     can_spill: bool,
     next_to_restore_partition_id: isize,
     spilled_small_partitions: Vec<Vec<usize>>,
-    spilled_merged_partitions: Vec<(SpilledData, bool, bool)>,
+    spilled_merged_partitions: Vec<(MergedPartition, bool, bool)>,
 }
 
 impl WindowPartitionBuffer {
@@ -110,7 +110,7 @@ impl WindowPartitionBuffer {
                 {
                     return self
                         .spiller
-                        .spill_with_partition(partition_id, DataBlock::concat(&data_blocks)?)
+                        .spill_with_partition(partition_id, data_blocks)
                         .await;
                 }
             }
@@ -127,8 +127,7 @@ impl WindowPartitionBuffer {
                     .partition_buffer
                     .fetch_data_blocks(partition_id, &option)?
                 {
-                    let data_block = DataBlock::concat(&data_blocks)?;
-                    partitions_to_spill.push((partition_id, data_block));
+                    partitions_to_spill.push((partition_id, data_blocks));
                     accumulated_bytes += partition_memory_size;
                 }
                 if accumulated_bytes >= spill_unit_size {
@@ -137,34 +136,20 @@ impl WindowPartitionBuffer {
             }
         }
 
-        if accumulated_bytes > 0 {
-            let spilled_data = self
-                .spiller
-                .spill_with_merged_partitions(partitions_to_spill)
-                .await?;
-            if let SpilledData::MergedPartition {
-                location,
-                partitions,
-            } = spilled_data
-            {
-                let index = self.spilled_merged_partitions.len();
-                for partition in partitions.iter() {
-                    self.spilled_small_partitions[partition.0].push(index);
-                }
-                self.spilled_merged_partitions.push((
-                    SpilledData::MergedPartition {
-                        location,
-                        partitions,
-                    },
-                    false,
-                    false,
-                ));
-                return Ok(());
-            }
+        if accumulated_bytes == 0 {
+            self.can_spill = false;
+            return Ok(());
         }
 
-        self.can_spill = false;
-
+        let spilled = self
+            .spiller
+            .spill_with_merged_partitions(partitions_to_spill)
+            .await?;
+        let index = self.spilled_merged_partitions.len();
+        for (id, _) in &spilled.partitions {
+            self.spilled_small_partitions[*id].push(index);
+        }
+        self.spilled_merged_partitions.push((spilled, false, false));
         Ok(())
     }
 
@@ -186,35 +171,31 @@ impl WindowPartitionBuffer {
                 if *restored {
                     continue;
                 }
-                if let SpilledData::MergedPartition {
+                let MergedPartition {
                     location,
                     partitions,
-                } = merged_partitions
-                {
-                    if out_of_memory_limit || *partial_restored {
-                        if let Some(pos) = partitions.iter().position(|p| p.0 == partition_id) {
-                            let data_range = &partitions[pos].1;
-                            let columns_layout = &partitions[pos].2;
-                            let data_block = self
-                                .spiller
-                                .read_range(location, data_range.clone(), columns_layout)
-                                .await?;
-                            self.restored_partition_buffer
-                                .add_data_block(partition_id, data_block);
-                            partitions.remove(pos);
-                            *partial_restored = true;
-                        }
-                    } else {
-                        let partitioned_data = self
+                } = merged_partitions;
+                if out_of_memory_limit || *partial_restored {
+                    if let Some(pos) = partitions.iter().position(|(id, _)| *id == partition_id) {
+                        let data_block = self
                             .spiller
-                            .read_merged_partitions(merged_partitions)
+                            .read_chunk(location, &partitions[pos].1)
                             .await?;
-                        for (partition_id, data_block) in partitioned_data.into_iter() {
-                            self.restored_partition_buffer
-                                .add_data_block(partition_id, data_block);
-                        }
-                        *restored = true;
+                        self.restored_partition_buffer
+                            .add_data_block(partition_id, data_block);
+                        partitions.remove(pos);
+                        *partial_restored = true;
+                    }
+                } else {
+                    let partitioned_data = self
+                        .spiller
+                        .read_merged_partitions(merged_partitions)
+                        .await?;
+                    for (partition_id, data_block) in partitioned_data.into_iter() {
+                        self.restored_partition_buffer
+                            .add_data_block(partition_id, data_block);
                     }
+                    *restored = true;
                 }
             }
 

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

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 mod partition_buffer;
+mod serialize;
 mod spiller;
 
 pub use partition_buffer::PartitionBuffer;