Merge branch 'main' into feat-gist-index

Author: youngsofun

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

@@ -13,8 +13,9 @@
 // limitations under the License.
 
 mod grace_join;
-mod grace_memory;
+pub mod grace_memory;
 mod grace_state;
 
 pub use grace_join::GraceHashJoin;
+pub use grace_memory::GraceMemoryJoin;
 pub use grace_state::*;

src/query/service/src/pipelines/processors/transforms/new_hash_join/hash_join_factory.rs

@@ -19,13 +19,14 @@ use std::sync::Mutex;
 use std::sync::PoisonError;
 use std::sync::Weak;
 
-use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_expression::FunctionContext;
 use databend_common_expression::HashMethodKind;
+use databend_common_pipeline_transforms::MemorySettings;
 use databend_common_sql::plans::JoinType;
 
+use crate::pipelines::memory_settings::MemorySettingsExt;
 use crate::pipelines::processors::HashJoinDesc;
 use crate::pipelines::processors::transforms::BasicHashJoinState;
 use crate::pipelines::processors::transforms::GraceHashJoin;
@@ -39,6 +40,9 @@ use crate::pipelines::processors::transforms::memory::SemiRightHashJoin;
 use crate::pipelines::processors::transforms::memory::left_join::OuterLeftHashJoin;
 use crate::pipelines::processors::transforms::new_hash_join::common::CStyleCell;
 use crate::pipelines::processors::transforms::new_hash_join::grace::GraceHashJoinState;
+use crate::pipelines::processors::transforms::new_hash_join::grace::GraceMemoryJoin;
+use crate::pipelines::processors::transforms::new_hash_join::hybrid::HybridHashJoin;
+use crate::pipelines::processors::transforms::new_hash_join::hybrid::HybridHashJoinState;
 use crate::sessions::QueryContext;
 
 pub struct HashJoinFactory {
@@ -49,6 +53,7 @@ pub struct HashJoinFactory {
     function_ctx: FunctionContext,
     grace_state: CStyleCell<HashMap<usize, Weak<GraceHashJoinState>>>,
     basic_state: CStyleCell<HashMap<usize, Weak<BasicHashJoinState>>>,
+    hybrid_state: CStyleCell<HashMap<usize, Weak<HybridHashJoinState>>>,
 }
 
 impl HashJoinFactory {
@@ -66,6 +71,7 @@ impl HashJoinFactory {
             mutex: Mutex::new(()),
             grace_state: CStyleCell::new(HashMap::new()),
             basic_state: CStyleCell::new(HashMap::new()),
+            hybrid_state: CStyleCell::new(HashMap::new()),
         })
     }
 
@@ -74,6 +80,7 @@ impl HashJoinFactory {
         let _locked = locked.unwrap_or_else(PoisonError::into_inner);
 
         let ctx = self.ctx.clone();
+
         match self.grace_state.as_mut().entry(id) {
             Entry::Occupied(v) => match v.get().upgrade() {
                 Some(v) => Ok(v),
@@ -110,6 +117,27 @@ impl HashJoinFactory {
         }
     }
 
+    pub fn create_hybrid_state(self: &Arc<Self>, level: usize) -> Result<Arc<HybridHashJoinState>> {
+        let locked = self.mutex.lock();
+        let _locked = locked.unwrap_or_else(PoisonError::into_inner);
+
+        match self.hybrid_state.as_mut().entry(level) {
+            Entry::Occupied(v) => match v.get().upgrade() {
+                Some(v) => Ok(v),
+                None => Err(ErrorCode::Internal(format!(
+                    "Error state: The level {} hybrid hash state has been destroyed.",
+                    level
+                ))),
+            },
+            Entry::Vacant(v) => {
+                let ctx = self.ctx.clone();
+                let hybrid_state = HybridHashJoinState::create(ctx, level, self.clone())?;
+                v.insert(Arc::downgrade(&hybrid_state));
+                Ok(hybrid_state)
+            }
+        }
+    }
+
     pub fn remove_basic_state(&self, id: usize) {
         let locked = self.mutex.lock();
         let _locked = locked.unwrap_or_else(PoisonError::into_inner);
@@ -124,64 +152,13 @@ impl HashJoinFactory {
     }
 
     pub fn create_hash_join(self: &Arc<Self>, typ: JoinType, id: usize) -> Result<Box<dyn Join>> {
-        let settings = self.ctx.get_settings();
+        // let settings = self.ctx.get_settings();
+        //
+        // if settings.get_force_join_data_spill()? {
+        //     return self.create_grace_join(typ, id);
+        // }
 
-        if settings.get_force_join_data_spill()? {
-            return self.create_grace_join(typ, id);
-        }
-
-        match typ {
-            JoinType::Inner => Ok(Box::new(InnerHashJoin::create(
-                &self.ctx,
-                self.function_ctx.clone(),
-                self.hash_method.clone(),
-                self.desc.clone(),
-                self.create_basic_state(id)?,
-            )?)),
-            JoinType::Left => Ok(Box::new(OuterLeftHashJoin::create(
-                &self.ctx,
-                self.function_ctx.clone(),
-                self.hash_method.clone(),
-                self.desc.clone(),
-                self.create_basic_state(id)?,
-            )?)),
-            JoinType::LeftAnti => Ok(Box::new(AntiLeftHashJoin::create(
-                &self.ctx,
-                self.function_ctx.clone(),
-                self.hash_method.clone(),
-                self.desc.clone(),
-                self.create_basic_state(id)?,
-            )?)),
-            JoinType::LeftSemi => Ok(Box::new(SemiLeftHashJoin::create(
-                &self.ctx,
-                self.function_ctx.clone(),
-                self.hash_method.clone(),
-                self.desc.clone(),
-                self.create_basic_state(id)?,
-            )?)),
-            JoinType::Right => Ok(Box::new(OuterRightHashJoin::create(
-                &self.ctx,
-                self.function_ctx.clone(),
-                self.hash_method.clone(),
-                self.desc.clone(),
-                self.create_basic_state(id)?,
-            )?)),
-            JoinType::RightSemi => Ok(Box::new(SemiRightHashJoin::create(
-                &self.ctx,
-                self.function_ctx.clone(),
-                self.hash_method.clone(),
-                self.desc.clone(),
-                self.create_basic_state(id)?,
-            )?)),
-            JoinType::RightAnti => Ok(Box::new(AntiRightHashJoin::create(
-                &self.ctx,
-                self.function_ctx.clone(),
-                self.hash_method.clone(),
-                self.desc.clone(),
-                self.create_basic_state(id)?,
-            )?)),
-            _ => unreachable!(),
-        }
+        Ok(Box::new(self.create_hybrid_join(typ, id)?))
     }
 
     pub fn create_grace_join(self: &Arc<Self>, typ: JoinType, id: usize) -> Result<Box<dyn Join>> {
@@ -322,4 +299,91 @@ impl HashJoinFactory {
             _ => unreachable!(),
         }
     }
+
+    /// Create a basic memory join (used internally by HybridHashJoin)
+    pub fn create_memory_join(
+        self: &Arc<Self>,
+        typ: JoinType,
+        level: usize,
+    ) -> Result<Box<dyn GraceMemoryJoin>> {
+        let basic_state = self.create_basic_state(level)?;
+        match typ {
+            JoinType::Inner => Ok(Box::new(InnerHashJoin::create(
+                &self.ctx,
+                self.function_ctx.clone(),
+                self.hash_method.clone(),
+                self.desc.clone(),
+                basic_state,
+            )?)),
+            JoinType::Left => Ok(Box::new(OuterLeftHashJoin::create(
+                &self.ctx,
+                self.function_ctx.clone(),
+                self.hash_method.clone(),
+                self.desc.clone(),
+                basic_state,
+            )?)),
+            JoinType::LeftAnti => Ok(Box::new(AntiLeftHashJoin::create(
+                &self.ctx,
+                self.function_ctx.clone(),
+                self.hash_method.clone(),
+                self.desc.clone(),
+                basic_state,
+            )?)),
+            JoinType::LeftSemi => Ok(Box::new(SemiLeftHashJoin::create(
+                &self.ctx,
+                self.function_ctx.clone(),
+                self.hash_method.clone(),
+                self.desc.clone(),
+                basic_state,
+            )?)),
+            JoinType::Right => Ok(Box::new(OuterRightHashJoin::create(
+                &self.ctx,
+                self.function_ctx.clone(),
+                self.hash_method.clone(),
+                self.desc.clone(),
+                basic_state,
+            )?)),
+            JoinType::RightSemi => Ok(Box::new(SemiRightHashJoin::create(
+                &self.ctx,
+                self.function_ctx.clone(),
+                self.hash_method.clone(),
+                self.desc.clone(),
+                basic_state,
+            )?)),
+            JoinType::RightAnti => Ok(Box::new(AntiRightHashJoin::create(
+                &self.ctx,
+                self.function_ctx.clone(),
+                self.hash_method.clone(),
+                self.desc.clone(),
+                basic_state,
+            )?)),
+            _ => unreachable!(),
+        }
+    }
+
+    /// Create a HybridHashJoin at the specified level
+    pub fn create_hybrid_join(
+        self: &Arc<Self>,
+        typ: JoinType,
+        level: usize,
+    ) -> Result<HybridHashJoin> {
+        let basic_state = self.create_basic_state(level)?;
+        let memory_join = self.create_memory_join(typ, level)?;
+        let memory_settings = MemorySettings::from_join_settings(&self.ctx)?;
+
+        // Use shared hybrid_state so all processors can coordinate spill
+        let hybrid_state = self.create_hybrid_state(level)?;
+
+        Ok(HybridHashJoin::create(
+            self.ctx.clone(),
+            self.function_ctx.clone(),
+            self.hash_method.clone(),
+            self.desc.clone(),
+            memory_settings,
+            hybrid_state,
+            basic_state,
+            memory_join,
+            typ,
+        ))
+    }
 }