optimize code

Author: xudong963

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

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::collections::HashMap;
+use std::collections::HashSet;
 
 use common_exception::Result;
 use common_functions::scalars::FunctionFactory;
@@ -28,15 +29,15 @@ use crate::sql::plans::JoinType;
 
 pub struct RightJoinDesc {
     /// Record rows in build side that are matched with rows in probe side.
-    pub(crate) build_indexes: RwLock<Vec<RowPtr>>,
+    pub(crate) build_indexes: RwLock<HashSet<RowPtr>>,
     /// Record row in build side that is matched how many rows in probe side.
     pub(crate) row_state: RwLock<HashMap<RowPtr, usize>>,
 }
 
 impl RightJoinDesc {
     pub fn create() -> Self {
         RightJoinDesc {
-            build_indexes: RwLock::new(vec![]),
+            build_indexes: RwLock::new(HashSet::new()),
             row_state: RwLock::new(HashMap::new()),
         }
     }

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

@@ -432,6 +432,37 @@ impl JoinHashTable {
             }
         }
     }
+
+    fn find_unmatched_build_indexes(&self) -> Result<Vec<RowPtr>> {
+        // For right join, build side will appear at lease once in the joined table
+        // Find the unmatched rows in build side
+        let mut unmatched_build_indexes = vec![];
+        {
+            let chunks = self.row_space.chunks.read().unwrap();
+            for (chunk_index, chunk) in chunks.iter().enumerate() {
+                for row_index in 0..chunk.num_rows() {
+                    let row_ptr = RowPtr {
+                        chunk_index: chunk_index as u32,
+                        row_index: row_index as u32,
+                        marker: None,
+                    };
+                    if !self
+                        .hash_join_desc
+                        .right_join_desc
+                        .build_indexes
+                        .read()
+                        .contains(&row_ptr)
+                    {
+                        let mut row_state = self.hash_join_desc.right_join_desc.row_state.write();
+                        row_state.entry(row_ptr).or_insert(0_usize);
+                        unmatched_build_indexes.push(row_ptr);
+                    }
+                }
+            }
+            drop(chunks);
+        }
+        Ok(unmatched_build_indexes)
+    }
 }
 
 #[async_trait::async_trait]
@@ -689,32 +720,7 @@ impl HashJoinState for JoinHashTable {
     }
 
     fn right_join_blocks(&self, blocks: &[DataBlock]) -> Result<Vec<DataBlock>> {
-        // For right join, build side will appear at lease once in the joined table
-        // Find the unmatched rows in build side
-        let mut unmatched_build_indexes = vec![];
-        {
-            for (chunk_index, chunk) in self.row_space.chunks.read().unwrap().iter().enumerate() {
-                for row_index in 0..chunk.num_rows() {
-                    let row_ptr = RowPtr {
-                        chunk_index: chunk_index as u32,
-                        row_index: row_index as u32,
-                        marker: None,
-                    };
-                    if !self
-                        .hash_join_desc
-                        .right_join_desc
-                        .build_indexes
-                        .read()
-                        .contains(&row_ptr)
-                    {
-                        let mut row_state = self.hash_join_desc.right_join_desc.row_state.write();
-                        row_state.entry(row_ptr).or_insert(0_usize);
-                        unmatched_build_indexes.push(row_ptr);
-                    }
-                }
-            }
-        }
-
+        let unmatched_build_indexes = self.find_unmatched_build_indexes()?;
         if unmatched_build_indexes.is_empty() && self.hash_join_desc.other_predicate.is_none() {
             return Ok(blocks.to_vec());
         }
@@ -770,7 +776,7 @@ impl HashJoinState for JoinHashTable {
         // If build_indexes size will greater build table size, we need filter the redundant rows for build side.
         let mut build_indexes = self.hash_join_desc.right_join_desc.build_indexes.write();
         let mut row_state = self.hash_join_desc.right_join_desc.row_state.write();
-        build_indexes.extend_from_slice(&unmatched_build_indexes);
+        build_indexes.extend(&unmatched_build_indexes);
         if build_indexes.len() > self.row_space.rows_number() {
             let mut bm = validity.into_mut().right().unwrap();
             Self::filter_rows_for_right_join(&mut bm, &build_indexes, &mut row_state);

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

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashSet;
 use std::iter::TrustedLen;
 
 use common_arrow::arrow::bitmap::Bitmap;
@@ -546,7 +547,7 @@ impl JoinHashTable {
                 {
                     let mut build_indexes =
                         self.hash_join_desc.right_join_desc.build_indexes.write();
-                    build_indexes.extend_from_slice(probe_result_ptrs);
+                    build_indexes.extend(probe_result_ptrs);
                     local_build_indexes.extend_from_slice(probe_result_ptrs);
                 }
                 for row_ptr in probe_result_ptrs.iter() {
@@ -655,7 +656,7 @@ impl JoinHashTable {
 
     pub(crate) fn filter_rows_for_right_join(
         bm: &mut MutableBitmap,
-        build_indexes: &[RowPtr],
+        build_indexes: &HashSet<RowPtr>,
         row_state: &mut std::collections::HashMap<RowPtr, usize>,
     ) {
         for (index, row) in build_indexes.iter().enumerate() {

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

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::hash::Hash;
+use std::hash::Hasher;
 use std::sync::RwLock;
 
 use common_datablocks::DataBlock;
@@ -36,7 +38,7 @@ impl Chunk {
     }
 }
 
-#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
+#[derive(Clone, Copy, Debug)]
 pub struct RowPtr {
     pub chunk_index: u32,
     pub row_index: u32,
@@ -110,3 +112,18 @@ impl RowSpace {
         }
     }
 }
+
+impl PartialEq for RowPtr {
+    fn eq(&self, other: &Self) -> bool {
+        self.chunk_index == other.chunk_index && self.row_index == other.row_index
+    }
+}
+
+impl Eq for RowPtr {}
+
+impl Hash for RowPtr {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.chunk_index.hash(state);
+        self.row_index.hash(state);
+    }
+}