x

Author: forsaken628

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

@@ -40,6 +40,7 @@ use crate::spillers::Spiller;
 use crate::spillers::SpillerConfig;
 use crate::spillers::SpillerDiskConfig;
 use crate::spillers::SpillerType;
+use crate::spillers::WindowSpiller;
 
 enum WindowBuffer {
     V1(WindowPartitionBuffer),
@@ -49,22 +50,23 @@ enum WindowBuffer {
 impl WindowBuffer {
     fn new(
         is_v2: bool,
-        spiller: Spiller,
+        partition_spiller: Spiller,
+        writer_spiller: WindowSpiller,
         num_partitions: usize,
         sort_block_size: usize,
         memory_settings: MemorySettings,
     ) -> Result<Self> {
         if is_v2 {
             let inner = WindowPartitionBufferV2::new(
-                spiller,
+                writer_spiller,
                 num_partitions,
                 sort_block_size,
                 memory_settings,
             )?;
             Ok(Self::V2(inner))
         } else {
             let inner = WindowPartitionBuffer::new(
-                spiller,
+                partition_spiller,
                 num_partitions,
                 sort_block_size,
                 memory_settings,
@@ -179,15 +181,18 @@ impl<S: DataProcessorStrategy> TransformWindowPartitionCollect<S> {
             use_parquet: settings.get_spilling_file_format()?.is_parquet(),
         };
 
-        // Create an inner `Spiller` to spill data.
+        // Create spillers for window operator.
         let operator = DataOperator::instance().spill_operator();
-        let spiller = Spiller::create(ctx, operator, spill_config)?;
+        let partition_spiller =
+            Spiller::create(ctx.clone(), operator.clone(), spill_config.clone())?;
+        let window_spiller = WindowSpiller::create(ctx, operator, spill_config)?;
 
         // Create the window partition buffer.
         let sort_block_size = settings.get_window_partition_sort_block_size()? as usize;
         let buffer = WindowBuffer::new(
             true,
-            spiller,
+            partition_spiller,
+            window_spiller,
             partitions.len(),
             sort_block_size,
             memory_settings,

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

@@ -22,7 +22,7 @@ use databend_common_pipeline_transforms::MemorySettings;
 use super::concat_data_blocks;
 use crate::spillers::SpillReader;
 use crate::spillers::SpillWriter;
-use crate::spillers::Spiller;
+use crate::spillers::WindowSpiller;
 
 #[async_trait::async_trait]
 pub trait Reader: Send {
@@ -46,7 +46,7 @@ pub trait Builder: Send + Sync {
 }
 
 #[async_trait::async_trait]
-impl Builder for Spiller {
+impl Builder for WindowSpiller {
     type Writer = SpillWriter;
 
     async fn create(&self, schema: Arc<DataSchema>) -> Result<SpillWriter> {
@@ -181,7 +181,7 @@ where
     }
 }
 
-pub(super) type WindowPartitionBufferV2 = PartitionBuffer<Spiller>;
+pub(super) type WindowPartitionBufferV2 = PartitionBuffer<WindowSpiller>;
 
 pub(super) struct PartitionBuffer<B>
 where B: Builder
@@ -196,9 +196,11 @@ where B: Builder
     next_to_restore_partition_id: isize,
 }
 
-impl PartitionBuffer<Spiller> {
+impl<B> PartitionBuffer<B>
+where B: Builder
+{
     pub fn new(
-        spiller: Spiller,
+        spiller: B,
         num_partitions: usize,
         sort_block_size: usize,
         memory_settings: MemorySettings,

src/query/service/src/spillers/adapter.rs

@@ -277,7 +277,7 @@ impl Spiller {
         };
 
         // Record statistics.
-        record_read_profile(location, &instant, data.len());
+        record_read_profile(location.is_local(), &instant, data.len());
 
         // Deserialize partitioned data block.
         let mut partitioned_data = Vec::with_capacity(partitions.len());
@@ -311,7 +311,7 @@ impl Spiller {
             Location::Remote(loc) => self.operator.read_with(loc).range(data_range).await?,
         };
 
-        record_read_profile(location, &instant, data.len());
+        record_read_profile(location.is_local(), &instant, data.len());
 
         deserialize_block(layout, data)
     }
@@ -354,20 +354,55 @@ impl Spiller {
             .cloned()
             .collect()
     }
+}
 
-    pub fn new_spill_writer(&self, schema: Arc<DataSchema>) -> Result<SpillWriter> {
+#[derive(Clone)]
+pub struct WindowWriterAdapter {
+    ctx: Arc<QueryContext>,
+    buffer_pool: Arc<BufferPool>,
+    chunk_size: usize,
+}
+
+impl WindowWriterAdapter {
+    fn add_spill_file(&self, location: Location, layout: Layout, size: usize) {
+        if location.is_remote() {
+            self.ctx.as_ref().incr_spill_progress(1, size);
+            self.ctx
+                .as_ref()
+                .add_spill_file(location.clone(), layout.clone());
+        }
+    }
+}
+
+pub type WindowSpiller = SpillerInner<WindowWriterAdapter>;
+
+impl WindowSpiller {
+    pub fn create(
+        ctx: Arc<QueryContext>,
+        operator: Operator,
+        config: SpillerConfig,
+    ) -> Result<Self> {
         let runtime = GlobalIORuntime::instance();
         let buffer_pool = BufferPool::create(
             runtime,
             WINDOW_SPILL_BUFFER_MEMORY_BYTES,
             WINDOW_SPILL_BUFFER_WORKERS,
         );
+        Self::new(
+            WindowWriterAdapter {
+                ctx,
+                buffer_pool,
+                chunk_size: WINDOW_SPILL_CHUNK_SIZE,
+            },
+            operator,
+            config,
+        )
+    }
 
+    pub fn new_spill_writer(&self, schema: Arc<DataSchema>) -> Result<SpillWriter> {
         Ok(SpillWriter {
             spiller: self.clone(),
-            buffer_pool,
-            dio: self.temp_dir.is_some(),
-            chunk_size: WINDOW_SPILL_CHUNK_SIZE,
+            chunk_size: self.adapter.chunk_size,
             schema,
             file_writer: None,
         })
@@ -389,9 +424,7 @@ const WINDOW_SPILL_BUFFER_WORKERS: usize = 2;
 const WINDOW_SPILL_CHUNK_SIZE: usize = 8 * 1024 * 1024;
 
 pub struct SpillWriter {
-    spiller: Spiller,
-    buffer_pool: Arc<BufferPool>,
-    dio: bool,
+    spiller: WindowSpiller,
     chunk_size: usize,
     schema: Arc<DataSchema>,
     file_writer: Option<FileWriter<UnionFileWriter>>,
@@ -405,7 +438,11 @@ impl SpillWriter {
 
         let writer = self
             .spiller
-            .new_file_writer(&self.schema, &self.buffer_pool, self.dio, self.chunk_size)
+            .new_file_writer(
+                &self.schema,
+                &self.spiller.adapter.buffer_pool,
+                self.chunk_size,
+            )
             .await?;
         self.file_writer = Some(writer);
         Ok(())
@@ -475,17 +512,15 @@ impl SpillWriter {
             schema: self.schema,
             parquet_metadata: Arc::new(metadata),
             union_file,
-            dio: self.dio,
         })
     }
 }
 
 pub struct SpillReader {
-    spiller: Spiller,
+    spiller: WindowSpiller,
     schema: Arc<DataSchema>,
     parquet_metadata: Arc<ParquetMetaData>,
     union_file: UnionFile,
-    dio: bool,
 }
 
 impl SpillReader {
@@ -500,7 +535,6 @@ impl SpillReader {
                 self.parquet_metadata.clone(),
                 &self.schema,
                 ordinals,
-                self.dio,
             )
             .await
     }
@@ -613,9 +647,7 @@ impl LiteSpiller {
                 Location::Local(_) => None,
             })
             .collect();
-        let op = self.0.local_operator.as_ref().unwrap_or(&self.0.operator);
-
-        op.delete_iter(files).await?;
+        self.0.operator.delete_iter(files).await?;
         Ok(())
     }
 }

src/query/service/src/spillers/inner.rs

@@ -100,7 +100,7 @@ pub trait SpillAdapter: Send + Sync + 'static {
 /// 3. Serialization and deserialization input data
 /// 4. Interact with the underlying storage engine to write and read spilled data
 #[derive(Clone)]
-pub struct SpillerInner<A: SpillAdapter> {
+pub struct SpillerInner<A> {
     pub(super) adapter: A,
     pub(super) operator: Operator,
     location_prefix: String,
@@ -111,7 +111,7 @@ pub struct SpillerInner<A: SpillAdapter> {
     _spiller_type: SpillerType,
 }
 
-impl<A: SpillAdapter> SpillerInner<A> {
+impl<A> SpillerInner<A> {
     pub fn new(adapter: A, operator: Operator, config: SpillerConfig) -> Result<Self> {
         let SpillerConfig {
             location_prefix,
@@ -139,17 +139,6 @@ impl<A: SpillAdapter> SpillerInner<A> {
         })
     }
 
-    /// Spill some [`DataBlock`] to storage. These blocks will be concat into one.
-    #[fastrace::trace(name = "Spiller::spill")]
-    pub async fn spill(&self, data_block: Vec<DataBlock>) -> Result<Location> {
-        let (location, layout, data_size) = self.spill_unmanage(data_block).await?;
-
-        // Record columns layout for spilled data.
-        self.adapter
-            .add_spill_file(location.clone(), layout, data_size);
-        Ok(location)
-    }
-
     async fn spill_unmanage(
         &self,
         data_block: Vec<DataBlock>,
@@ -179,13 +168,6 @@ impl<A: SpillAdapter> SpillerInner<A> {
         format!("{}/{}", self.location_prefix, GlobalUniqName::unique())
     }
 
-    /// Read a certain file to a [`DataBlock`].
-    #[fastrace::trace(name = "Spiller::read_spilled_file")]
-    pub async fn read_spilled_file(&self, location: &Location) -> Result<DataBlock> {
-        let layout = self.adapter.get_spill_layout(location).unwrap();
-        self.read_unmanage_spilled_file(location, &layout).await
-    }
-
     async fn read_unmanage_spilled_file(
         &self,
         location: &Location,
@@ -219,7 +201,7 @@ impl<A: SpillAdapter> SpillerInner<A> {
             Location::Remote(loc) => self.operator.read(loc).await?,
         };
 
-        record_read_profile(location, &instant, data.len());
+        record_read_profile(location.is_local(), &instant, data.len());
 
         deserialize_block(columns_layout, data)
     }
@@ -273,6 +255,26 @@ impl<A: SpillAdapter> SpillerInner<A> {
     }
 }
 
+impl<A: SpillAdapter> SpillerInner<A> {
+    /// Spill some [`DataBlock`] to storage. These blocks will be concat into one.
+    #[fastrace::trace(name = "Spiller::spill")]
+    pub async fn spill(&self, data_block: Vec<DataBlock>) -> Result<Location> {
+        let (location, layout, data_size) = self.spill_unmanage(data_block).await?;
+
+        // Record columns layout for spilled data.
+        self.adapter
+            .add_spill_file(location.clone(), layout, data_size);
+        Ok(location)
+    }
+
+    /// Read a certain file to a [`DataBlock`].
+    #[fastrace::trace(name = "Spiller::read_spilled_file")]
+    pub async fn read_spilled_file(&self, location: &Location) -> Result<DataBlock> {
+        let layout = self.adapter.get_spill_layout(location).unwrap();
+        self.read_unmanage_spilled_file(location, &layout).await
+    }
+}
+
 pub(super) fn record_write_profile(is_local: bool, start: &Instant, write_bytes: usize) {
     if !is_local {
         Profile::record_usize_profile(ProfileStatisticsName::RemoteSpillWriteCount, 1);
@@ -291,23 +293,20 @@ pub(super) fn record_write_profile(is_local: bool, start: &Instant, write_bytes:
     }
 }
 
-pub(super) fn record_read_profile(location: &Location, start: &Instant, read_bytes: usize) {
-    match location {
-        Location::Remote(_) => {
-            Profile::record_usize_profile(ProfileStatisticsName::RemoteSpillReadCount, 1);
-            Profile::record_usize_profile(ProfileStatisticsName::RemoteSpillReadBytes, read_bytes);
-            Profile::record_usize_profile(
-                ProfileStatisticsName::RemoteSpillReadTime,
-                start.elapsed().as_millis() as usize,
-            );
-        }
-        Location::Local(_) => {
-            Profile::record_usize_profile(ProfileStatisticsName::LocalSpillReadCount, 1);
-            Profile::record_usize_profile(ProfileStatisticsName::LocalSpillReadBytes, read_bytes);
-            Profile::record_usize_profile(
-                ProfileStatisticsName::LocalSpillReadTime,
-                start.elapsed().as_millis() as usize,
-            );
-        }
+pub(super) fn record_read_profile(is_local: bool, start: &Instant, read_bytes: usize) {
+    if is_local {
+        Profile::record_usize_profile(ProfileStatisticsName::RemoteSpillReadCount, 1);
+        Profile::record_usize_profile(ProfileStatisticsName::RemoteSpillReadBytes, read_bytes);
+        Profile::record_usize_profile(
+            ProfileStatisticsName::RemoteSpillReadTime,
+            start.elapsed().as_millis() as usize,
+        );
+    } else {
+        Profile::record_usize_profile(ProfileStatisticsName::LocalSpillReadCount, 1);
+        Profile::record_usize_profile(ProfileStatisticsName::LocalSpillReadBytes, read_bytes);
+        Profile::record_usize_profile(
+            ProfileStatisticsName::LocalSpillReadTime,
+            start.elapsed().as_millis() as usize,
+        );
     }
 }