add TransformRightJoin

Author: xudong963

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

@@ -54,6 +54,7 @@ pub use transforms::KeyU64HashTable;
 pub use transforms::KeyU8HashTable;
 pub use transforms::MarkJoinCompactor;
 pub use transforms::ProjectionTransform;
+pub use transforms::RightJoinCompactor;
 pub use transforms::SerializerHashTable;
 pub use transforms::SinkBuildHashTable;
 pub use transforms::SortMergeCompactor;

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

@@ -12,17 +12,36 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
+
 use common_exception::Result;
 use common_functions::scalars::FunctionFactory;
 use parking_lot::RwLock;
 
 use crate::evaluator::EvalNode;
 use crate::evaluator::Evaluator;
+use crate::pipelines::processors::transforms::hash_join::row::RowPtr;
 use crate::pipelines::processors::transforms::hash_join::MarkJoinDesc;
 use crate::sql::executor::HashJoin;
 use crate::sql::executor::PhysicalScalar;
 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>>,
+    /// 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![]),
+            row_state: RwLock::new(HashMap::new()),
+        }
+    }
+}
+
 pub struct HashJoinDesc {
     pub(crate) build_keys: Vec<EvalNode>,
     pub(crate) probe_keys: Vec<EvalNode>,
@@ -31,6 +50,7 @@ pub struct HashJoinDesc {
     pub(crate) marker_join_desc: MarkJoinDesc,
     /// Whether the Join are derived from correlated subquery.
     pub(crate) from_correlated_subquery: bool,
+    pub(crate) right_join_desc: RightJoinDesc,
 }
 
 impl HashJoinDesc {
@@ -50,6 +70,7 @@ impl HashJoinDesc {
                 marker_index: join.marker_index,
             },
             from_correlated_subquery: join.from_correlated_subquery,
+            right_join_desc: RightJoinDesc::create(),
         })
     }
 

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

@@ -44,4 +44,7 @@ pub trait HashJoinState: Send + Sync {
 
     /// Get mark join results
     fn mark_join_blocks(&self) -> Result<Vec<DataBlock>>;
+
+    /// Get right join results
+    fn right_join_blocks(&self, blocks: &[DataBlock]) -> Result<Vec<DataBlock>>;
 }

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

@@ -17,6 +17,7 @@ use std::fmt::Debug;
 use std::sync::Arc;
 use std::sync::Mutex;
 
+use common_arrow::arrow::bitmap::Bitmap;
 use common_arrow::arrow::bitmap::MutableBitmap;
 use common_base::base::tokio::sync::Notify;
 use common_datablocks::DataBlock;
@@ -35,7 +36,9 @@ use common_datavalues::DataField;
 use common_datavalues::DataSchema;
 use common_datavalues::DataSchemaRef;
 use common_datavalues::DataSchemaRefExt;
+use common_datavalues::DataType;
 use common_datavalues::DataTypeImpl;
+use common_datavalues::DataValue;
 use common_datavalues::NullableType;
 use common_exception::ErrorCode;
 use common_exception::Result;
@@ -684,4 +687,94 @@ impl HashJoinState for JoinHashTable {
         let build_block = self.row_space.gather(&row_ptrs)?;
         Ok(vec![self.merge_eq_block(&marker_block, &build_block)?])
     }
+
+    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![];
+        {
+            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)
+                    {
+                        unmatched_build_indexes.push(row_ptr);
+                    }
+                }
+            }
+        }
+
+        let unmatched_build_block = self.row_space.gather(&unmatched_build_indexes)?;
+        // Create null block for unmatched rows in probe side
+        let null_probe_block = DataBlock::create(
+            self.probe_schema.clone(),
+            self.probe_schema
+                .fields()
+                .iter()
+                .map(|df| {
+                    df.data_type()
+                        .clone()
+                        .create_constant_column(&DataValue::Null, unmatched_build_indexes.len())
+                })
+                .collect::<Result<Vec<_>>>()?,
+        );
+        let mut merged_block = self.merge_eq_block(&null_probe_block, &unmatched_build_block)?;
+        merged_block = DataBlock::concat_blocks(&[blocks, &[merged_block]].concat())?;
+
+        if self.hash_join_desc.other_predicate.is_none() {
+            return Ok(vec![merged_block]);
+        }
+
+        let (bm, all_true, all_false) = self.get_other_filters(
+            &merged_block,
+            self.hash_join_desc.other_predicate.as_ref().unwrap(),
+        )?;
+
+        if all_true {
+            return Ok(vec![merged_block]);
+        }
+
+        let validity = match (bm, all_false) {
+            (Some(b), _) => b,
+            (None, true) => Bitmap::new_zeroed(merged_block.num_rows()),
+            // must be one of above
+            _ => unreachable!(),
+        };
+
+        let probe_column_len = self.probe_schema.fields().len();
+        let probe_columns = merged_block.columns()[0..probe_column_len]
+            .iter()
+            .map(|c| Self::set_validity(c, &validity))
+            .collect::<Result<Vec<_>>>()?;
+        let probe_block = DataBlock::create(self.probe_schema.clone(), probe_columns);
+        let build_block = DataBlock::create(
+            self.row_space.data_schema.clone(),
+            merged_block.columns()[probe_column_len..].to_vec(),
+        );
+        merged_block = self.merge_eq_block(&probe_block, &build_block)?;
+
+        // If there are only non-equi conditions, build_indexes size will greater build table size
+        // Because the case will cause cross join.
+        // We need filter the redundant rows for build side.
+        let build_indexes = self.hash_join_desc.right_join_desc.build_indexes.read();
+        let mut row_state = self.hash_join_desc.right_join_desc.row_state.write();
+        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);
+            let predicate = BooleanColumn::from_arrow_data(bm.into()).arc();
+            return Ok(vec![DataBlock::filter_block(merged_block, &predicate)?]);
+        }
+
+        Ok(vec![merged_block])
+    }
 }

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

@@ -159,15 +159,8 @@ impl JoinHashTable {
                 }
             }
             JoinType::Right => {
-                if self.hash_join_desc.other_predicate.is_none() {
-                    let result =
-                        self.right_join::<false, _, _>(hash_table, probe_state, keys_iter, input)?;
-                    return Ok(vec![result]);
-                } else {
-                    let result =
-                        self.right_join::<true, _, _>(hash_table, probe_state, keys_iter, input)?;
-                    return Ok(vec![result]);
-                }
+                let result = self.right_join::<_, _>(hash_table, probe_state, keys_iter, input)?;
+                return Ok(vec![result]);
             }
             Mark => {
                 results.push(DataBlock::empty());
@@ -531,7 +524,7 @@ impl JoinHashTable {
         DataBlock::filter_block(merged_block, &predicate)
     }
 
-    fn right_join<const WITH_OTHER_CONJUNCT: bool, Key, IT>(
+    fn right_join<Key, IT>(
         &self,
         hash_table: &HashMap<Key, Vec<RowPtr>>,
         probe_state: &mut ProbeState,
@@ -542,102 +535,45 @@ impl JoinHashTable {
         Key: HashTableKeyable + Clone + 'static,
         IT: Iterator<Item = Key> + TrustedLen,
     {
-        let build_indexes = &mut probe_state.build_indexs;
+        let local_build_indexes = &mut probe_state.build_indexs;
         let probe_indexes = &mut probe_state.probe_indexs;
         let valids = &probe_state.valids;
         let mut validity = MutableBitmap::new();
-        let mut row_state = std::collections::HashMap::new();
         for (i, key) in keys_iter.enumerate() {
             let probe_result_ptr = Self::probe_key(hash_table, key, valids, i);
             if let Some(v) = probe_result_ptr {
                 let probe_result_ptrs = v.get_value();
-                build_indexes.extend_from_slice(probe_result_ptrs);
+                {
+                    let mut build_indexes =
+                        self.hash_join_desc.right_join_desc.build_indexes.write();
+                    build_indexes.extend_from_slice(probe_result_ptrs);
+                    local_build_indexes.extend_from_slice(probe_result_ptrs);
+                }
                 for row_ptr in probe_result_ptrs.iter() {
-                    row_state
-                        .entry(*row_ptr)
-                        .and_modify(|e| *e += 1)
-                        .or_insert(1_usize);
+                    {
+                        let mut row_state = self.hash_join_desc.right_join_desc.row_state.write();
+                        row_state
+                            .entry(*row_ptr)
+                            .and_modify(|e| *e += 1)
+                            .or_insert(1_usize);
+                    }
                 }
                 probe_indexes.extend(std::iter::repeat(i as u32).take(probe_result_ptrs.len()));
                 validity.extend_constant(probe_result_ptrs.len(), true);
             }
         }
 
-        // 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 kv in hash_table.iter() {
-            for v in kv.get_value() {
-                if !build_indexes.contains(v) {
-                    unmatched_build_indexes.push(*v);
-                }
-            }
-        }
-        build_indexes.extend_from_slice(&unmatched_build_indexes);
-        let build_block = self.row_space.gather(build_indexes)?;
+        let build_block = self.row_space.gather(local_build_indexes)?;
         let probe_block = DataBlock::block_take_by_indices(input, probe_indexes)?;
         let validity: Bitmap = validity.into();
         let nullable_columns = probe_block
             .columns()
             .iter()
             .map(|c| Self::set_validity(c, &validity))
             .collect::<Result<Vec<_>>>()?;
-        let mut nullable_probe_block =
-            DataBlock::create(self.probe_schema.clone(), nullable_columns);
-        // Create null block for unmatched rows in probe side
-        let null_probe_block = DataBlock::create(
-            self.probe_schema.clone(),
-            nullable_probe_block
-                .columns()
-                .iter()
-                .map(|c| {
-                    c.data_type()
-                        .create_constant_column(&DataValue::Null, unmatched_build_indexes.len())
-                })
-                .collect::<Result<Vec<_>>>()?,
-        );
-
-        nullable_probe_block = DataBlock::concat_blocks(&[nullable_probe_block, null_probe_block])?;
-        let mut merged_block = self.merge_eq_block(&nullable_probe_block, &build_block)?;
-        if !WITH_OTHER_CONJUNCT {
-            return Ok(merged_block);
-        }
-
-        let (bm, all_true, all_false) = self.get_other_filters(
-            &merged_block,
-            self.hash_join_desc.other_predicate.as_ref().unwrap(),
-        )?;
-
-        if all_true {
-            return Ok(merged_block);
-        }
-
-        let validity = match (bm, all_false) {
-            (Some(b), _) => b,
-            (None, true) => Bitmap::new_zeroed(merged_block.num_rows()),
-            // must be one of above
-            _ => unreachable!(),
-        };
-
-        let nullable_columns = nullable_probe_block
-            .columns()
-            .iter()
-            .map(|c| Self::set_validity(c, &validity))
-            .collect::<Result<Vec<_>>>()?;
-        nullable_probe_block = DataBlock::create(self.probe_schema.clone(), nullable_columns);
-        merged_block = self.merge_eq_block(&nullable_probe_block, &build_block)?;
-
-        // If there are only non-equi conditions, build_indexes size will greater build table size
-        // Because the case will cause cross join.
-        // We need filter the redundant rows for build side.
-        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);
-            let predicate = BooleanColumn::from_arrow_data(bm.into()).arc();
-            return DataBlock::filter_block(merged_block, &predicate);
-        }
+        let nullable_probe_block = DataBlock::create(self.probe_schema.clone(), nullable_columns);
 
-        Ok(merged_block)
+        self.merge_eq_block(&nullable_probe_block, &build_block)
     }
 
     // modify the bm by the value row_state
@@ -717,7 +653,7 @@ impl JoinHashTable {
         }
     }
 
-    fn filter_rows_for_right_join(
+    pub(crate) fn filter_rows_for_right_join(
         bm: &mut MutableBitmap,
         build_indexes: &[RowPtr],
         row_state: &mut std::collections::HashMap<RowPtr, usize>,
@@ -737,7 +673,7 @@ impl JoinHashTable {
     }
 
     // return an (option bitmap, all_true, all_false)
-    fn get_other_filters(
+    pub(crate) fn get_other_filters(
         &self,
         merged_block: &DataBlock,
         filter: &EvalNode,

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

@@ -38,6 +38,7 @@ mod transform_window_func;
 
 pub mod group_by;
 mod transform_merge_block;
+mod transform_right_join;
 
 pub use aggregator::AggregatorParams;
 pub use aggregator::AggregatorTransformParams;
@@ -79,6 +80,8 @@ pub use transform_mark_join::TransformMarkJoin;
 pub use transform_merge_block::TransformMergeBlock;
 pub use transform_project::TransformProject;
 pub use transform_rename::TransformRename;
+pub use transform_right_join::RightJoinCompactor;
+pub use transform_right_join::TransformRightJoin;
 pub use transform_sort_merge::SortMergeCompactor;
 pub use transform_sort_merge::TransformSortMerge;
 pub use transform_sort_partial::get_sort_descriptions;