Skip to content

Commit 2df2a3e

Browse files
zhang2014zhang2014
authored
refactor(query): restructure hash join memory implementations (databendlabs#19199)
* refactor(query): restructure hash join memory implementations * z * z * z * z * z * z * z * z * z * z * z * z * z * z * z * z
1 parent 39abacc commit 2df2a3e

23 files changed

Lines changed: 2057 additions & 78 deletions

‎src/common/base/src/base/barrier.rs‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
// Copyright 2021 Datafuse Labs
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
use std::sync::Mutex;
16+
use std::sync::PoisonError;
17+
18+
use tokio::sync::watch;
19+
20+
#[derive(Debug)]
21+
struct BarrierState {
22+
waker: watch::Sender<usize>,
23+
arrived: usize,
24+
generation: usize,
25+
26+
n: usize,
27+
}
28+
29+
pub struct Barrier {
30+
state: Mutex<BarrierState>,
31+
wait: watch::Receiver<usize>,
32+
}
33+
34+
impl Barrier {
35+
pub fn new(mut n: usize) -> Barrier {
36+
let (waker, wait) = watch::channel(0);
37+
38+
if n == 0 {
39+
n = 1;
40+
}
41+
42+
Barrier {
43+
state: Mutex::new(BarrierState {
44+
n,
45+
waker,
46+
arrived: 0,
47+
generation: 1,
48+
}),
49+
wait,
50+
}
51+
}
52+
53+
pub async fn wait(&self) -> BarrierWaitResult {
54+
let (generation, is_leader) = {
55+
let locked = self.state.lock();
56+
let mut state = locked.unwrap_or_else(PoisonError::into_inner);
57+
58+
let is_leader = state.arrived == 0;
59+
let generation = state.generation;
60+
state.arrived += 1;
61+
62+
if state.arrived == state.n {
63+
state
64+
.waker
65+
.send(state.generation)
66+
.expect("there is at least one receiver");
67+
state.arrived = 0;
68+
state.generation += 1;
69+
return BarrierWaitResult(is_leader);
70+
}
71+
72+
(generation, is_leader)
73+
};
74+
75+
let mut wait = self.wait.clone();
76+
77+
loop {
78+
let _ = wait.changed().await;
79+
80+
if *wait.borrow() >= generation {
81+
break;
82+
}
83+
}
84+
85+
BarrierWaitResult(is_leader)
86+
}
87+
88+
pub fn reduce_quorum(&self, n: usize) {
89+
let locked = self.state.lock();
90+
let mut state = locked.unwrap_or_else(PoisonError::into_inner);
91+
state.n -= n;
92+
93+
if state.arrived >= state.n {
94+
state
95+
.waker
96+
.send(state.generation)
97+
.expect("there is at least one receiver");
98+
state.arrived = 0;
99+
state.generation += 1;
100+
}
101+
}
102+
}
103+
104+
#[derive(Debug, Clone)]
105+
pub struct BarrierWaitResult(bool);
106+
107+
impl BarrierWaitResult {
108+
pub fn is_leader(&self) -> bool {
109+
self.0
110+
}
111+
}

‎src/common/base/src/base/mod.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15+
mod barrier;
1516
mod build_info;
1617
mod dma;
1718
mod drop_callback;
@@ -30,6 +31,7 @@ mod take_mut;
3031
mod uniq_id;
3132
mod watch_notify;
3233

34+
pub use barrier::Barrier;
3335
pub use build_info::*;
3436
pub use dma::*;
3537
pub use drop_callback::DropCallback;

‎src/query/service/src/physical_plans/physical_cache_scan.rs‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -85,9 +85,19 @@ impl IPhysicalPlan for CacheScan {
8585
max_block_size,
8686
))
8787
}
88-
Some(HashJoinStateRef::NewHashJoinState(hash_join_state)) => {
88+
Some(HashJoinStateRef::NewHashJoinState(hash_join_state, column_map)) => {
89+
let mut column_offsets = Vec::with_capacity(column_indexes.len());
90+
for index in column_indexes {
91+
let Some(offset) = column_map.get(index) else {
92+
return Err(ErrorCode::Internal(format!(
93+
"Hash join cache column {} not found in build projection",
94+
index
95+
)));
96+
};
97+
column_offsets.push(*offset);
98+
}
8999
CacheSourceState::NewHashJoinCacheState(NewHashJoinCacheState::new(
90-
column_indexes.clone(),
100+
column_offsets,
91101
hash_join_state.clone(),
92102
))
93103
}

‎src/query/service/src/physical_plans/physical_hash_join.rs‎

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ use crate::pipelines::processors::transforms::RuntimeFiltersDesc;
6666
use crate::pipelines::processors::transforms::TransformHashJoin;
6767
use crate::pipelines::processors::transforms::TransformHashJoinBuild;
6868
use crate::pipelines::processors::transforms::TransformHashJoinProbe;
69+
use crate::sessions::QueryContext;
6970

7071
// Type aliases to simplify complex return types
7172
type JoinConditionsResult = (
@@ -270,19 +271,29 @@ impl IPhysicalPlan for HashJoin {
270271
let (enable_optimization, _) = builder.merge_into_get_optimization_flag(self);
271272

272273
if desc.single_to_inner.is_none()
273-
&& (self.join_type == JoinType::Inner || self.join_type == JoinType::Left)
274+
&& matches!(
275+
self.join_type,
276+
JoinType::Inner
277+
| JoinType::Left
278+
| JoinType::LeftSemi
279+
| JoinType::LeftAnti
280+
| JoinType::Right
281+
| JoinType::RightSemi
282+
| JoinType::RightAnti
283+
)
274284
&& experimental_new_join
275285
&& !enable_optimization
286+
&& !self.need_hold_hash_table
276287
{
277288
return self.build_new_join_pipeline(builder, desc);
278289
}
279290

280291
// Create the join state with optimization flags
281292
let state = self.build_state(builder)?;
282293

283-
if let Some((build_cache_index, _)) = self.build_side_cache_info {
294+
if let Some((build_cache_index, _)) = &self.build_side_cache_info {
284295
builder.hash_join_states.insert(
285-
build_cache_index,
296+
*build_cache_index,
286297
HashJoinStateRef::OldHashJoinState(state.clone()),
287298
);
288299
}
@@ -413,15 +424,18 @@ impl HashJoin {
413424
{
414425
let state = factory.create_basic_state(0)?;
415426

416-
if let Some((build_cache_index, _)) = self.build_side_cache_info {
427+
if let Some((build_cache_index, column_map)) = &self.build_side_cache_info {
417428
builder.hash_join_states.insert(
418-
build_cache_index,
419-
HashJoinStateRef::NewHashJoinState(state.clone()),
429+
*build_cache_index,
430+
HashJoinStateRef::NewHashJoinState(state.clone(), column_map.clone()),
420431
);
421432
}
422433
}
423434

435+
let mut sub_query_ctx = QueryContext::create_from(&builder.ctx);
436+
std::mem::swap(&mut builder.ctx, &mut sub_query_ctx);
424437
self.build.build_pipeline(builder)?;
438+
std::mem::swap(&mut builder.ctx, &mut sub_query_ctx);
425439
let mut build_sinks = builder.main_pipeline.take_sinks();
426440

427441
self.probe.build_pipeline(builder)?;
@@ -440,7 +454,8 @@ impl HashJoin {
440454

441455
debug_assert_eq!(build_sinks.len(), probe_sinks.len());
442456

443-
let stage_sync_barrier = Arc::new(Barrier::new(output_len));
457+
let barrier = databend_common_base::base::Barrier::new(output_len);
458+
let stage_sync_barrier = Arc::new(barrier);
444459
let mut join_sinks = Vec::with_capacity(output_len * 2);
445460
let mut join_pipe_items = Vec::with_capacity(output_len);
446461
for (build_sink, probe_sink) in build_sinks.into_iter().zip(probe_sinks.into_iter()) {

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ use databend_common_pipeline::core::ExecutionInfo;
2323
use databend_common_pipeline::core::Pipeline;
2424
use databend_common_pipeline::core::always_callback;
2525
use databend_common_settings::Settings;
26+
use databend_common_sql::IndexType;
2627

2728
use super::PipelineBuilderData;
2829
use crate::interpreters::CreateTableInterpreter;
@@ -38,7 +39,7 @@ use crate::sessions::QueryContext;
3839
#[derive(Clone)]
3940
pub enum HashJoinStateRef {
4041
OldHashJoinState(Arc<HashJoinState>),
41-
NewHashJoinState(Arc<BasicHashJoinState>),
42+
NewHashJoinState(Arc<BasicHashJoinState>, HashMap<IndexType, usize>),
4243
}
4344

4445
pub struct PipelineBuilder {

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,3 +198,14 @@ pub fn wrap_true_validity(
198198
NullableColumn::new_column(col, validity).into()
199199
}
200200
}
201+
202+
pub fn wrap_nullable_block(input: &DataBlock) -> DataBlock {
203+
let input_num_rows = input.num_rows();
204+
let true_validity = Bitmap::new_constant(true, input_num_rows);
205+
let nullable_columns = input
206+
.columns()
207+
.iter()
208+
.map(|c| wrap_true_validity(c, input_num_rows, &true_validity))
209+
.collect::<Vec<_>>();
210+
DataBlock::new(nullable_columns, input_num_rows)
211+
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ pub struct HashJoinDesc {
6363
pub(crate) probe_projections: ColumnSet,
6464
pub(crate) probe_to_build: Vec<(usize, (bool, bool))>,
6565
pub(crate) build_schema: DataSchemaRef,
66+
pub(crate) probe_schema: DataSchemaRef,
6667
}
6768

6869
#[derive(Debug, Clone)]
@@ -138,6 +139,7 @@ impl HashJoinDesc {
138139
build_projection: join.build_projections.clone(),
139140
probe_projections: join.probe_projections.clone(),
140141
build_schema: join.build.output_schema()?,
142+
probe_schema: join.probe.output_schema()?,
141143
})
142144
}
143145

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ mod transform_hash_join_build;
3030
mod transform_hash_join_probe;
3131
mod util;
3232

33+
pub use common::wrap_nullable_block;
3334
pub use common::wrap_true_validity;
3435
pub use desc::HashJoinDesc;
3536
pub use desc::RuntimeFilterDesc;

‎src/query/service/src/pipelines/processors/transforms/new_hash_join/grace/grace_memory.rs‎

Lines changed: 52 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,12 @@ use crate::pipelines::processors::transforms::BasicHashJoinState;
1818
use crate::pipelines::processors::transforms::HashJoinHashTable;
1919
use crate::pipelines::processors::transforms::InnerHashJoin;
2020
use crate::pipelines::processors::transforms::Join;
21-
use crate::pipelines::processors::transforms::memory::outer_left_join::OuterLeftHashJoin;
21+
use crate::pipelines::processors::transforms::memory::AntiLeftHashJoin;
22+
use crate::pipelines::processors::transforms::memory::AntiRightHashJoin;
23+
use crate::pipelines::processors::transforms::memory::OuterRightHashJoin;
24+
use crate::pipelines::processors::transforms::memory::SemiLeftHashJoin;
25+
use crate::pipelines::processors::transforms::memory::SemiRightHashJoin;
26+
use crate::pipelines::processors::transforms::memory::left_join::OuterLeftHashJoin;
2227

2328
pub trait GraceMemoryJoin: Join {
2429
fn reset_memory(&mut self);
@@ -52,6 +57,14 @@ fn reset_basic_state(state: &BasicHashJoinState) {
5257
state.build_queue.as_mut().clear();
5358
}
5459

60+
if !state.scan_map.is_empty() {
61+
state.scan_map.as_mut().clear();
62+
}
63+
64+
if !state.scan_queue.is_empty() {
65+
state.scan_queue.as_mut().clear();
66+
}
67+
5568
*state.hash_table.as_mut() = HashJoinHashTable::Null;
5669
}
5770

@@ -68,3 +81,41 @@ impl GraceMemoryJoin for OuterLeftHashJoin {
6881
reset_basic_state(&self.basic_state);
6982
}
7083
}
84+
85+
impl GraceMemoryJoin for SemiLeftHashJoin {
86+
fn reset_memory(&mut self) {
87+
self.performance_context.clear();
88+
reset_basic_state(&self.basic_state);
89+
}
90+
}
91+
92+
impl GraceMemoryJoin for AntiLeftHashJoin {
93+
fn reset_memory(&mut self) {
94+
self.performance_context.clear();
95+
reset_basic_state(&self.basic_state);
96+
}
97+
}
98+
99+
impl GraceMemoryJoin for OuterRightHashJoin {
100+
fn reset_memory(&mut self) {
101+
self.finished = false;
102+
self.performance_context.clear();
103+
reset_basic_state(&self.basic_state);
104+
}
105+
}
106+
107+
impl GraceMemoryJoin for SemiRightHashJoin {
108+
fn reset_memory(&mut self) {
109+
self.finished = false;
110+
self.performance_context.clear();
111+
reset_basic_state(&self.basic_state);
112+
}
113+
}
114+
115+
impl GraceMemoryJoin for AntiRightHashJoin {
116+
fn reset_memory(&mut self) {
117+
self.finished = false;
118+
self.performance_context.clear();
119+
reset_basic_state(&self.basic_state);
120+
}
121+
}

0 commit comments

Comments
 (0)