Merge pull request #10711 from xudong963/left_join_oom

fix: left join oom

Author: xudong963

src/query/service/src/pipelines/pipeline_builder.rs

@@ -829,7 +829,7 @@ impl PipelineBuilder {
             || join.join_type == JoinType::Single)
             && join.non_equi_conditions.is_empty()
         {
-            self.main_pipeline.resize(1)?;
+            // self.main_pipeline.resize(1)?;
             self.main_pipeline.add_transform(|input, output| {
                 let transform = TransformLeftJoin::try_create(
                     input,

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

@@ -309,20 +309,22 @@ impl JoinHashTable {
     }
 
     pub(crate) fn rest_block(&self) -> Result<DataBlock> {
-        let rest_probe_blocks = self.hash_join_desc.join_state.rest_probe_blocks.read();
-        if rest_probe_blocks.is_empty() {
+        let mut rest_pairs = self.hash_join_desc.join_state.rest_pairs.write();
+        if rest_pairs.0.is_empty() {
             return Ok(DataBlock::empty());
         }
-        let probe_block = DataBlock::concat(&rest_probe_blocks)?;
-        let rest_build_indexes = self.hash_join_desc.join_state.rest_build_indexes.read();
-        let mut build_block = self.row_space.gather(&rest_build_indexes)?;
+        let probe_block = DataBlock::concat(&rest_pairs.0)?;
+        rest_pairs.0.clear();
+        let mut build_block = self.row_space.gather(&rest_pairs.1)?;
+        rest_pairs.1.clear();
         // For left join, wrap nullable for build block
         if matches!(
             self.hash_join_desc.join_type,
             JoinType::Left | JoinType::Single | JoinType::Full
         ) {
-            let validity = self.hash_join_desc.join_state.validity.read();
-            let validity: Bitmap = (*validity).clone().into();
+            let mut validity_state = self.hash_join_desc.join_state.validity.write();
+            let validity: Bitmap = (*validity_state).clone().into();
+            validity_state.clear();
             let num_rows = validity.len();
             let nullable_columns = if self.row_space.datablocks().is_empty() {
                 build_block
@@ -342,7 +344,6 @@ impl JoinHashTable {
             };
             build_block = DataBlock::new(nullable_columns, num_rows);
         }
-
         self.merge_eq_block(&build_block, &probe_block)
     }
 

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

@@ -43,18 +43,17 @@ pub struct JoinState {
     /// Record rows in build side that are matched with rows in probe side.
     /// It's order-sensitive, aligned with the order of rows in merged block.
     pub(crate) build_indexes: RwLock<Vec<RowPtr>>,
-    pub(crate) rest_build_indexes: RwLock<Vec<RowPtr>>,
-    pub(crate) rest_probe_blocks: RwLock<Vec<DataBlock>>,
+    /// Rest build indexes and probe blocks
+    pub(crate) rest_pairs: RwLock<(Vec<DataBlock>, Vec<RowPtr>)>,
     pub(crate) validity: RwLock<MutableBitmap>,
 }
 
 impl JoinState {
     pub fn create() -> Result<Self> {
         Ok(JoinState {
             build_indexes: RwLock::new(Vec::with_capacity(JOIN_MAX_BLOCK_SIZE)),
-            rest_build_indexes: RwLock::new(Vec::with_capacity(JOIN_MAX_BLOCK_SIZE)),
-            rest_probe_blocks: RwLock::new(Vec::with_capacity(JOIN_MAX_BLOCK_SIZE)),
             validity: RwLock::new(MutableBitmap::with_capacity(JOIN_MAX_BLOCK_SIZE)),
+            rest_pairs: Default::default(),
         })
     }
 }

src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/left_join.rs

@@ -243,10 +243,9 @@ impl JoinHashTable {
         }
 
         if !WITH_OTHER_CONJUNCT {
-            let mut rest_build_indexes = self.hash_join_desc.join_state.rest_build_indexes.write();
-            rest_build_indexes.extend(local_build_indexes);
-            let mut rest_probe_blocks = self.hash_join_desc.join_state.rest_probe_blocks.write();
-            rest_probe_blocks.push(probe_block);
+            let mut rest_pairs = self.hash_join_desc.join_state.rest_pairs.write();
+            rest_pairs.1.extend(local_build_indexes);
+            rest_pairs.0.push(probe_block);
             let validity: Bitmap = validity.into();
             let mut validity_state = self.hash_join_desc.join_state.validity.write();
             validity_state.extend_from_bitmap(&validity);

src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/right_join.rs

@@ -136,10 +136,9 @@ impl JoinHashTable {
             probe_block = DataBlock::new(nullable_columns, validity.len());
         }
 
-        let mut rest_build_indexes = self.hash_join_desc.join_state.rest_build_indexes.write();
-        let mut rest_probe_blocks = self.hash_join_desc.join_state.rest_probe_blocks.write();
-        rest_probe_blocks.push(probe_block);
-        rest_build_indexes.extend(local_build_indexes);
+        let mut rest_pairs = self.hash_join_desc.join_state.rest_pairs.write();
+        rest_pairs.0.push(probe_block);
+        rest_pairs.1.extend(local_build_indexes);
 
         Ok(probed_blocks)
     }

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

@@ -40,6 +40,17 @@ impl Compactor for LeftJoinCompactor {
         self.hash_join_state.interrupt();
     }
 
+    fn use_partial_compact() -> bool {
+        true
+    }
+
+    fn compact_partial(&mut self, blocks: &mut Vec<DataBlock>) -> Result<Vec<DataBlock>> {
+        let res = self.hash_join_state.left_join_blocks(blocks)?;
+        // Original blocks are already in res, so clear them.
+        blocks.clear();
+        Ok(res)
+    }
+
     // `compact_final` is called when all the blocks are pushed
     fn compact_final(&self, blocks: &[DataBlock]) -> Result<Vec<DataBlock>> {
         self.hash_join_state.left_join_blocks(blocks)