risingwave_batch_executors/executor/join/
lookup_join_base.rs

1// Copyright 2024 RisingWave 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
15use std::marker::PhantomData;
16
17use futures::StreamExt;
18use futures_async_stream::try_stream;
19use itertools::Itertools;
20use prometheus::core::Atomic;
21use risingwave_common::array::DataChunk;
22use risingwave_common::bitmap::FilterByBitmap;
23use risingwave_common::catalog::Schema;
24use risingwave_common::hash::{HashKey, NullBitmap, PrecomputedBuildHasher};
25use risingwave_common::memory::MemoryContext;
26use risingwave_common::metrics::TrAdderAtomic;
27use risingwave_common::row::Row;
28use risingwave_common::types::{DataType, ToOwnedDatum};
29use risingwave_common::util::chunk_coalesce::DataChunkBuilder;
30use risingwave_common::util::sort_util::{OrderType, cmp_datum_iter};
31use risingwave_common_estimate_size::EstimateSize;
32use risingwave_expr::expr::BoxedExpression;
33
34use super::AsOfDesc;
35use crate::error::BatchError;
36use crate::executor::join::chunked_data::ChunkedData;
37use crate::executor::{
38    BoxedDataChunkListStream, BoxedExecutor, BufferChunkExecutor, EquiJoinParams, HashJoinExecutor,
39    JoinHashMap, JoinType, LookupExecutorBuilder, RowId, utils,
40};
41use crate::task::ShutdownToken;
42
43/// Lookup Join Base.
44/// Used by `LocalLookupJoinExecutor` and `DistributedLookupJoinExecutor`.
45pub struct LookupJoinBase<K, B: LookupExecutorBuilder> {
46    pub join_type: JoinType,
47    pub condition: Option<BoxedExpression>,
48    pub outer_side_input: BoxedExecutor,
49    pub outer_side_data_types: Vec<DataType>, // Data types of all columns of outer side table
50    pub outer_side_key_idxs: Vec<usize>,
51    pub inner_side_builder: B,
52    pub inner_side_key_types: Vec<DataType>, // Data types only of key columns of inner side table
53    pub inner_side_key_idxs: Vec<usize>,
54    pub null_safe: Vec<bool>,
55    pub lookup_prefix_len: usize,
56    pub chunk_builder: DataChunkBuilder,
57    pub schema: Schema,
58    pub output_indices: Vec<usize>,
59    pub chunk_size: usize,
60    pub asof_desc: Option<AsOfDesc>,
61    pub identity: String,
62    pub shutdown_rx: ShutdownToken,
63    pub mem_ctx: MemoryContext,
64    pub _phantom: PhantomData<K>,
65}
66
67const AT_LEAST_OUTER_SIDE_ROWS: usize = 512;
68
69impl<K: HashKey, B: LookupExecutorBuilder> LookupJoinBase<K, B> {
70    /// High level Execution flow:
71    /// Repeat 1-3:
72    ///   1. Read N rows from outer side input and send keys to inner side builder after
73    ///      deduplication.
74    ///   2. Inner side input lookups inner side table with keys and builds hash map.
75    ///   3. Outer side rows join each inner side rows by probing the hash map.
76    #[try_stream(boxed, ok = DataChunk, error = BatchError)]
77    pub async fn do_execute(mut self: Box<Self>) {
78        let outer_side_schema = self.outer_side_input.schema().clone();
79
80        let null_matched = K::Bitmap::from_bool_vec(self.null_safe);
81
82        let mut outer_side_batch_read_stream: BoxedDataChunkListStream =
83            utils::batch_read(self.outer_side_input.execute(), AT_LEAST_OUTER_SIDE_ROWS);
84
85        while let Some(chunk_list) = outer_side_batch_read_stream.next().await {
86            let chunk_list = chunk_list?;
87
88            // Group rows with the same key datums together
89            let groups = chunk_list
90                .iter()
91                .flat_map(|chunk| {
92                    chunk.rows().map(|row| {
93                        self.outer_side_key_idxs
94                            .iter()
95                            .take(self.lookup_prefix_len)
96                            .map(|&idx| row.datum_at(idx).to_owned_datum())
97                            .collect_vec()
98                    })
99                })
100                .sorted_by(|a, b| cmp_datum_iter(a, b, std::iter::repeat(OrderType::default())))
101                .dedup()
102                .collect_vec();
103
104            self.inner_side_builder.reset();
105            for row_key in groups {
106                self.inner_side_builder.add_scan_range(row_key).await?;
107            }
108            let inner_side_input = self.inner_side_builder.build_executor().await?;
109
110            // Lookup join outer side will become the probe side of hash join,
111            // while its inner side will become the build side of hash join.
112            let hash_join_probe_side_input = Box::new(BufferChunkExecutor::new(
113                outer_side_schema.clone(),
114                chunk_list,
115            ));
116            let hash_join_build_side_input = inner_side_input;
117            let hash_join_probe_data_types = self.outer_side_data_types.clone();
118            let hash_join_build_data_types = hash_join_build_side_input.schema().data_types();
119            let hash_join_probe_side_key_idxs = self.outer_side_key_idxs.clone();
120            let hash_join_build_side_key_idxs = self.inner_side_key_idxs.clone();
121
122            let full_data_types = [
123                hash_join_probe_data_types.clone(),
124                hash_join_build_data_types.clone(),
125            ]
126            .concat();
127
128            // This round's private context releases remaining charges on completion, errors, or
129            // cancellation, before the next round starts.
130            let mem_ctx = MemoryContext::new(Some(self.mem_ctx.clone()), TrAdderAtomic::new(0));
131
132            let mut build_side = Vec::new_in(mem_ctx.global_allocator());
133            let mut build_row_count = 0;
134            #[for_await]
135            for build_chunk in hash_join_build_side_input.execute() {
136                let build_chunk = build_chunk?;
137                if build_chunk.cardinality() > 0 {
138                    build_row_count += build_chunk.cardinality();
139                    let chunk_estimated_heap_size = build_chunk.estimated_heap_size() as i64;
140                    build_side.push(build_chunk);
141                    mem_ctx.add_unchecked(chunk_estimated_heap_size);
142                }
143            }
144            let mut hash_map = JoinHashMap::with_capacity_and_hasher_in(
145                build_row_count,
146                PrecomputedBuildHasher,
147                mem_ctx.global_allocator(),
148            );
149            let mut next_build_row_with_same_key =
150                ChunkedData::with_chunk_sizes(build_side.iter().map(|c| c.capacity()))?;
151
152            // Build hash map
153            for (build_chunk_id, build_chunk) in build_side.iter().enumerate() {
154                let build_keys = K::build_many(&hash_join_build_side_key_idxs, build_chunk);
155
156                for (build_row_id, build_key) in build_keys
157                    .into_iter()
158                    .enumerate()
159                    .filter_by_bitmap(build_chunk.visibility())
160                {
161                    // Only insert key to hash map if it is consistent with the null safe
162                    // restriction.
163                    if build_key.null_bitmap().is_subset(&null_matched) {
164                        let row_id = RowId::new(build_chunk_id, build_row_id);
165                        let build_key_estimated_heap_size = build_key.estimated_heap_size() as i64;
166                        next_build_row_with_same_key[row_id] = hash_map.insert(build_key, row_id);
167                        mem_ctx.add_unchecked(build_key_estimated_heap_size);
168                    }
169                }
170            }
171
172            let params = EquiJoinParams::new(
173                hash_join_probe_side_input,
174                hash_join_probe_data_types,
175                hash_join_probe_side_key_idxs,
176                build_side,
177                hash_join_build_data_types,
178                full_data_types,
179                hash_map,
180                next_build_row_with_same_key,
181                self.chunk_size,
182                self.shutdown_rx.clone(),
183                self.asof_desc.clone(),
184            );
185
186            assert!(
187                !matches!(
188                    self.join_type,
189                    JoinType::RightOuter
190                        | JoinType::RightSemi
191                        | JoinType::RightAnti
192                        | JoinType::FullOuter
193                ),
194                "lookup join does not support {:?}",
195                self.join_type
196            );
197            #[for_await]
198            for chunk in HashJoinExecutor::probe_and_project(
199                params,
200                self.join_type,
201                self.condition.as_ref(),
202                &self.output_indices,
203                self.schema.data_types(),
204            ) {
205                yield chunk?
206            }
207        }
208    }
209}
risingwave_batch_executors/executor/join/lookup_join_base.rs

risingwave_batch_executors/executor/join/
lookup_join_base.rs
