risingwave_frontend/optimizer/delta_join_solver.rs
1// Copyright 2022 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
15//! The solver for delta join, which determines lookup order of a join plan.
16//! All collection types in this module should be `BTree` to ensure determinism between runs.
17//!
18//! # Representation of Multi-Way Join
19//!
20//! In this module, `a*` means lookup executor that looks-up the arrangement `a` of the current
21//! epoch. `a` means looks-up arrangement `a` of the previous epoch.
22//!
23//! Delta joins only support inner equal join. The solver is based on the following formula (take
24//! 3-way join as an example):
25//!
26//! ```plain
27//! d((A join1 B) join2 C)
28//! = ((A + dA) join1 (B + dB)) join2 (C + dC) - A join1 B join2 C
29//! = (A join1 B + A join1 dB + dA join1 (B + dB)) join2 (C + dC) - A join1 B join2 C
30//! = A join1 B join2 (C + dC) + A join1 dB join2 (C + dC) + dA join1 (B + dB) join2 (C + dC) - A join1 B join2 C
31//! = A join1 B join2 dC + A join1 dB join2 (C + dC) + dA join1 (B + dB) join2 (C + dC)
32//! = dA join1 (B + dB) join2 (C + dC) + dB join1 A join2 (C + dC) + dC join2 B join1 A
33//!
34//! join1 means A join B using condition #1,
35//! join2 means B join C using condition #2.
36//! ```
37//!
38//! Inner joins satisfy commutative law and associative laws, so we can switch them back and forth
39//! between joins.
40//!
41//! ... which generates the following look-up graph:
42//!
43//! ```plain
44//! a -> b* -> c* -> output3
45//! b -> a -> c* -> output2
46//! c -> b -> a -> output1
47//! ```
48//!
49//! And the final output is `output1 <concat> output2 <concat> output3`. The concatenation is
50//! ordered: all items from output 1 must appear before output 2.
51//!
52//! TODO: support dynamic filter and filter condition in lookups.
53//!
54//! # What need the caller do?
55//!
56//! After solving the delta join plan, caller will need to do a lot of things.
57//!
58//! * Use the correct index for every stream input. By looking at the first lookup fragment in each
59//! row, we can decide whether to use `a.x` or `a.y` as index for stream input.
60//! * Insert exchanges between lookups of different distribution. Generally, if the whole row is
61//! operating on the same join key, we only need to do 1v1 exchanges between lookups. However, it
62//! would be possible that a row of lookup first join `a.x == b.x`, then `a.y == c.y`. In this
63//! case, we will need to insert hash exchange between these two lookups.
64//! * Ensure the order of union. Always union from the last row to the first row.
65//! * Insert exchange before union. Still the case for `a.x == b.x`, then `a.y == c.y`, it is
66//! possible that every lookup path produces different distribution. We need to shuffle them
67//! before feeding data to union.
68
69#![expect(dead_code)]
70
71use std::collections::{BTreeMap, BTreeSet};
72
73use anyhow::{Result, anyhow};
74use itertools::Itertools;
75
76#[derive(Debug, Clone, Copy, Eq, PartialEq, PartialOrd, Ord)]
77pub struct JoinTable(pub usize);
78
79/// Represents whether `left` and `right` can be joined using a condition.
80#[derive(Debug, Clone, Eq, PartialEq)]
81pub struct JoinEdge {
82 pub left: JoinTable,
83 pub right: JoinTable,
84 pub left_join_key: Vec<usize>,
85 pub right_join_key: Vec<usize>,
86}
87
88impl JoinEdge {
89 /// Reverse the order of the edge.
90 pub fn reverse(&self) -> Self {
91 Self {
92 left: self.right,
93 right: self.left,
94 left_join_key: self.right_join_key.clone(),
95 right_join_key: self.left_join_key.clone(),
96 }
97 }
98}
99
100/// Decides how to place arrangements over lookup nodes. Given a 3-way join example:
101///
102/// ```plain
103/// a -> 1* -> 2* ->
104/// b -> 3 -> 4* ->
105/// c -> 5 -> 6 ->
106/// ```
107///
108/// If user provides the multi-way join order of `(a join b) join c`, and set the strategy to be
109/// [`ArrangeStrategy::LeftFirst`], and if three tables are of the same join key (the graph is
110/// fully-connected and no shuffle needed), then we will place the arrangements over the lookup
111/// nodes in the following way:
112///
113/// ```plain
114/// a -> b* -> c* ->
115/// b -> a -> c* ->
116/// c -> a -> b ->
117/// ```
118///
119/// The left side of joins will be preferred for the left lookups when selecting arrangements.
120///
121/// If strategy is set to [`ArrangeStrategy::RightFirst`],
122///
123/// ```plain
124/// a -> c* -> b* ->
125/// b -> c -> a* ->
126/// c -> b -> a ->
127/// ```
128///
129/// The right side of joins will be preferred for the left lookups when selecting arrangements.
130#[derive(Clone, Debug, PartialEq, Eq)]
131pub enum ArrangeStrategy {
132 /// The left-most table will be preferred to be the first lookup table.
133 LeftFirst,
134 /// The right-most table will be preferred to be the first lookup table.
135 RightFirst,
136}
137
138/// Decides how to place stream inputs. Given a 3-way join example:
139///
140/// ```plain
141/// x -> 1* -> 2* ->
142/// x -> 3 -> 4* ->
143/// x -> 5 -> 6 ->
144/// ```
145///
146/// ... where `n*` means lookup this epoch. If user provides the multi-way join order of `(a join b)
147/// join c`, and set the strategy to be [`StreamStrategy::LeftThisEpoch`], then we will place the
148/// stream side as:
149///
150/// ```plain
151/// a -> 1* -> 2* ->
152/// b -> 3 -> 4* ->
153/// c -> 5 -> 6 ->
154/// ```
155///
156/// ... where `a`, the left-most table in multi-way join, passes through most number of
157/// lookup-this-epoch.
158///
159/// If strategy is set to [`StreamStrategy::RightThisEpoch`],
160///
161/// ```plain
162/// c -> 1* -> 2* ->
163/// b -> 3 -> 4* ->
164/// a -> 5 -> 6 ->
165/// ```
166///
167/// ... where `c`, the right-most table in multi-way join, passes through most number of
168/// lookup-this-epoch.
169///
170/// More lookup-this-epochs in a lookup row mean more latency when a barrier is flushed. If lookup
171/// executor is set to lookup this epoch, it will wait until a barrier from the arrangement side
172/// before actually starting lookups and joins.
173#[derive(Clone, Debug, PartialEq, Eq)]
174pub enum StreamStrategy {
175 /// The left-most table will be preferred to be the stream to lookup this epoch.
176 LeftThisEpoch,
177 /// The right-most table will be preferred to be the stream to lookup this epoch.
178 RightThisEpoch,
179}
180
181/// Given a multi-way inner join plan, the solver will produces the most efficient way of getting
182/// those join done.
183///
184/// ## Input
185///
186/// [`DeltaJoinSolver`] needs the information of the cost of joining two tables. You'll need to
187/// provide edges of the join graph. Each [`JoinEdge`] `(a, b) -> cost` describes the cost of inner
188/// join a and b.
189///
190/// ## Output
191///
192/// The `solve` function will return a vector of [`LookupPath`]. See [`LookupPath`] for more.
193pub struct DeltaJoinSolver {
194 /// Strategy of arrangement placement, see docs of [`ArrangeStrategy`] for more details.
195 arrange_strategy: ArrangeStrategy,
196
197 /// Strategy of stream placement, see docs of [`StreamStrategy`] for more details.
198 stream_strategy: StreamStrategy,
199
200 /// Possible combination of joins. If `(a, b)` is in `edges`, then `a` and `b` can be joined.
201 /// The edge is bi-directional, so only one pair (either `a, b` or `b, a`) needs to be present
202 /// in this vector.
203 edges: Vec<JoinEdge>,
204
205 /// The recommended join order from optimizer in a multi-way join plan node.
206 join_order: Vec<JoinTable>,
207}
208
209/// A row in the lookup plan, which includes the stream side arrangement, and all arrangements to be
210/// placed in the lookup node.
211///
212/// For example, if we have `LookupPath(a, vec![b, c])`, then we have a row:
213///
214/// ```plain
215/// a -> b -> c
216/// ```
217///
218/// In `Vec<LookupPath>`, the first row always contains the most lookup-this-epochs. For example, if
219/// we have:
220///
221/// ```plain
222/// vec![
223/// LookupPath(a, vec![b, c]),
224/// LookupPath(b, vec![c, a]),
225/// LookupPath(c, vec![a, b])
226/// ]
227/// ```
228///
229/// Then it means...
230///
231/// ```plain
232/// a -> b* -> c* ->
233/// b -> c -> a* ->
234/// c -> a -> b ->
235/// ```
236#[derive(Clone, Debug, Eq, PartialEq)]
237pub struct LookupPath(JoinTable, pub Vec<JoinTable>);
238
239struct SolverEnv {
240 /// Stores all join edges in map. [`JoinEdge`]'s left side is always the key in map.
241 join_edge: BTreeMap<JoinTable, Vec<JoinEdge>>,
242
243 /// Placement order of arrangements
244 arrange_placement_order: Vec<JoinTable>,
245
246 /// Placement order of streams
247 stream_placement_order: Vec<JoinTable>,
248}
249
250impl SolverEnv {
251 fn build_from(solver: &DeltaJoinSolver) -> Self {
252 let mut join_edge = BTreeMap::new();
253
254 for table in &solver.join_order {
255 join_edge.insert(*table, vec![]);
256 }
257 for edge in &solver.edges {
258 join_edge.get_mut(&edge.left).unwrap().push(edge.clone());
259 join_edge.get_mut(&edge.right).unwrap().push(edge.reverse());
260 }
261
262 Self {
263 arrange_placement_order: match solver.arrange_strategy {
264 ArrangeStrategy::LeftFirst => solver.join_order.clone(),
265 ArrangeStrategy::RightFirst => {
266 solver.join_order.iter().copied().rev().collect_vec()
267 }
268 },
269 stream_placement_order: match solver.stream_strategy {
270 StreamStrategy::LeftThisEpoch => solver.join_order.clone(),
271 StreamStrategy::RightThisEpoch => {
272 solver.join_order.iter().copied().rev().collect_vec()
273 }
274 },
275 join_edge,
276 }
277 }
278}
279
280impl DeltaJoinSolver {
281 /// Generate a lookup path using the user provided strategy. The lookup path is generated in the
282 /// following way:
283 ///
284 /// * Firstly, we find all tables that can be joined with the current join table set. The tables
285 /// should have the same join key as the current join set.
286 /// * If there are multiple tables that satisfy this condition, we will pick tables according to
287 /// [`ArrangeStrategy`].
288 /// * If not, then we have to do a shuffle. We pick all tables that have join condition with the
289 /// current table set.
290 /// * And then pick a table using [`ArrangeStrategy`].
291 ///
292 /// Basically, this algorithm will greedily pick all tables with the same join key first (so
293 /// that there won't be shuffle). Then, switch a distribution, and do this process again.
294 fn find_lookup_path(
295 &self,
296 solver_env: &SolverEnv,
297 input_stream: JoinTable,
298 ) -> Result<LookupPath> {
299 // The table available to query
300 let mut current_table_set = BTreeSet::new();
301 current_table_set.insert(input_stream);
302
303 // The distribution of the current lookup.
304 let mut current_distribution = vec![];
305
306 // The final lookup path.
307 let mut path = vec![];
308
309 fn satisfies_distribution(
310 current_distribution: &[(JoinTable, Vec<usize>)],
311 edge: &JoinEdge,
312 ) -> bool {
313 // TODO: if `current_distribution` is a `BTreeSet`, we can know if the distribution in
314 // O(1). But as we generally have few tables, so we do a linear scan on tables.
315 if current_distribution.is_empty() {
316 return true;
317 }
318
319 for (table, join_key) in current_distribution {
320 if table == &edge.left && join_key == &edge.left_join_key {
321 return true;
322 }
323 }
324
325 false
326 }
327
328 'next_table: loop {
329 assert!(
330 path.len() < self.join_order.len(),
331 "internal error: infinite loop"
332 );
333
334 // step 1: find tables that can be joined with `current_table_set` and satisfy the
335 // current distribution.
336 let mut reachable_tables = BTreeMap::new();
337
338 for current_table in ¤t_table_set {
339 for edge in &solver_env.join_edge[current_table] {
340 if !current_table_set.contains(&edge.right)
341 && satisfies_distribution(¤t_distribution, edge)
342 {
343 reachable_tables.insert(edge.right, edge.clone());
344 }
345 }
346 }
347
348 // step 2: place arrangements according to the arrange strategy, update current
349 // distribution and current table set.
350 for table in &solver_env.arrange_placement_order {
351 if let Some(edge) = reachable_tables.get(table) {
352 current_table_set.insert(edge.right);
353 path.push(edge.right);
354 current_distribution.push((edge.right, edge.right_join_key.clone()));
355 continue 'next_table;
356 }
357 }
358
359 // no table can be joined using the same distribution at this point.
360
361 // step 3: find all tables that can be joined with `current_table_set`, regardless of
362 // their distribution.
363 let mut reachable_tables = BTreeMap::new();
364
365 for current_table in ¤t_table_set {
366 for edge in &solver_env.join_edge[current_table] {
367 if !current_table_set.contains(&edge.right) {
368 reachable_tables.insert(edge.right, edge.clone());
369 }
370 }
371 }
372
373 // step 4: place arrangements according to the arrange strategy, update current
374 // distribution and current table set.
375 for table in &solver_env.arrange_placement_order {
376 if let Some(edge) = reachable_tables.get(table) {
377 current_table_set.insert(edge.right);
378 path.push(edge.right);
379 current_distribution.clear();
380 current_distribution.push((edge.right, edge.right_join_key.clone()));
381 continue 'next_table;
382 }
383 }
384
385 // step 5: no tables can be joined any more, what happened?
386 if self.join_order.len() - 1 == path.len() {
387 break;
388 } else {
389 // no table can be joined, while path is still incomplete
390 return Err(anyhow!(
391 "join plan cannot be generated, tables not connected."
392 ));
393 }
394 }
395
396 Ok(LookupPath(input_stream, path))
397 }
398
399 pub fn solve(&self) -> Result<Vec<LookupPath>> {
400 let solver_env = SolverEnv::build_from(self);
401
402 solver_env
403 .stream_placement_order
404 .iter()
405 .map(|x| self.find_lookup_path(&solver_env, *x))
406 .collect()
407 }
408
409 pub fn new(
410 stream_strategy: StreamStrategy,
411 edges: Vec<JoinEdge>,
412 join_order: Vec<JoinTable>,
413 ) -> Self {
414 Self {
415 stream_strategy,
416 edges,
417 join_order,
418 arrange_strategy: ArrangeStrategy::LeftFirst,
419 }
420 }
421}
422
423#[cfg(test)]
424mod tests {
425 use super::*;
426
427 #[test]
428 fn test_2way_join() {
429 let solver = DeltaJoinSolver {
430 arrange_strategy: ArrangeStrategy::LeftFirst,
431 stream_strategy: StreamStrategy::LeftThisEpoch,
432 edges: vec![JoinEdge {
433 left: JoinTable(1),
434 right: JoinTable(2),
435 left_join_key: vec![2, 3],
436 right_join_key: vec![3, 2],
437 }],
438 join_order: vec![JoinTable(1), JoinTable(2)],
439 };
440
441 let result = solver.solve().unwrap();
442
443 assert_eq!(
444 result,
445 vec![
446 LookupPath(JoinTable(1), vec![JoinTable(2)]),
447 LookupPath(JoinTable(2), vec![JoinTable(1)]),
448 ]
449 );
450
451 let solver = DeltaJoinSolver {
452 arrange_strategy: ArrangeStrategy::RightFirst,
453 stream_strategy: StreamStrategy::LeftThisEpoch,
454 edges: vec![JoinEdge {
455 left: JoinTable(1),
456 right: JoinTable(2),
457 left_join_key: vec![2, 3],
458 right_join_key: vec![3, 2],
459 }],
460 join_order: vec![JoinTable(1), JoinTable(2)],
461 };
462
463 let result = solver.solve().unwrap();
464
465 assert_eq!(
466 result,
467 vec![
468 LookupPath(JoinTable(1), vec![JoinTable(2)]),
469 LookupPath(JoinTable(2), vec![JoinTable(1)]),
470 ]
471 );
472
473 let solver = DeltaJoinSolver {
474 arrange_strategy: ArrangeStrategy::LeftFirst,
475 stream_strategy: StreamStrategy::RightThisEpoch,
476 edges: vec![JoinEdge {
477 left: JoinTable(1),
478 right: JoinTable(2),
479 left_join_key: vec![2, 3],
480 right_join_key: vec![3, 2],
481 }],
482 join_order: vec![JoinTable(1), JoinTable(2)],
483 };
484
485 let result = solver.solve().unwrap();
486
487 assert_eq!(
488 result,
489 vec![
490 LookupPath(JoinTable(2), vec![JoinTable(1)]),
491 LookupPath(JoinTable(1), vec![JoinTable(2)]),
492 ]
493 );
494
495 let solver = DeltaJoinSolver {
496 arrange_strategy: ArrangeStrategy::RightFirst,
497 stream_strategy: StreamStrategy::RightThisEpoch,
498 edges: vec![JoinEdge {
499 left: JoinTable(1),
500 right: JoinTable(2),
501 left_join_key: vec![2, 3],
502 right_join_key: vec![3, 2],
503 }],
504 join_order: vec![JoinTable(1), JoinTable(2)],
505 };
506
507 let result = solver.solve().unwrap();
508
509 assert_eq!(
510 result,
511 vec![
512 LookupPath(JoinTable(2), vec![JoinTable(1)]),
513 LookupPath(JoinTable(1), vec![JoinTable(2)]),
514 ]
515 );
516 }
517
518 #[test]
519 fn test_3way_join_one_key() {
520 // Table 1: [2, 3] (composite key x)
521 // Table 2: [3, 2] (composite key x)
522 // Table 3: [1, 1] (composite key x)
523 // t1.x == t2.x == t3.x
524
525 let solver = DeltaJoinSolver {
526 arrange_strategy: ArrangeStrategy::LeftFirst,
527 stream_strategy: StreamStrategy::LeftThisEpoch,
528 edges: vec![
529 JoinEdge {
530 left: JoinTable(1),
531 right: JoinTable(2),
532 left_join_key: vec![2, 3],
533 right_join_key: vec![3, 2],
534 },
535 JoinEdge {
536 left: JoinTable(2),
537 right: JoinTable(3),
538 left_join_key: vec![3, 2],
539 right_join_key: vec![1, 1],
540 },
541 JoinEdge {
542 left: JoinTable(3),
543 right: JoinTable(1),
544 left_join_key: vec![1, 1],
545 right_join_key: vec![2, 3],
546 },
547 ],
548 join_order: vec![JoinTable(1), JoinTable(2), JoinTable(3)],
549 };
550
551 let result = solver.solve().unwrap();
552
553 assert_eq!(
554 result,
555 vec![
556 LookupPath(JoinTable(1), vec![JoinTable(2), JoinTable(3)]),
557 LookupPath(JoinTable(2), vec![JoinTable(1), JoinTable(3)]),
558 LookupPath(JoinTable(3), vec![JoinTable(1), JoinTable(2)])
559 ]
560 );
561 }
562
563 #[test]
564 fn test_3way_join_two_keys() {
565 // Table 1: [0] (key x)
566 // Table 2: [0] (key x), [1] (key y)
567 // Table 3: [1] (key y)
568 // t1.x == t2.x and t2.y == t3.y
569 //
570 // t1 cannot directly join with t3 in this case, as they don't share the same key
571
572 let solver = DeltaJoinSolver {
573 arrange_strategy: ArrangeStrategy::LeftFirst,
574 stream_strategy: StreamStrategy::LeftThisEpoch,
575 edges: vec![
576 JoinEdge {
577 left: JoinTable(1),
578 right: JoinTable(2),
579 left_join_key: vec![0],
580 right_join_key: vec![0],
581 },
582 JoinEdge {
583 left: JoinTable(2),
584 right: JoinTable(3),
585 left_join_key: vec![1],
586 right_join_key: vec![1],
587 },
588 ],
589 join_order: vec![JoinTable(1), JoinTable(2), JoinTable(3)],
590 };
591
592 let result = solver.solve().unwrap();
593
594 assert_eq!(
595 result,
596 vec![
597 LookupPath(JoinTable(1), vec![JoinTable(2), JoinTable(3)]),
598 LookupPath(JoinTable(2), vec![JoinTable(1), JoinTable(3)]),
599 LookupPath(JoinTable(3), vec![JoinTable(2), JoinTable(1)])
600 ]
601 );
602 }
603
604 #[test]
605 fn test_invalid_plan() {
606 let solver = DeltaJoinSolver {
607 arrange_strategy: ArrangeStrategy::LeftFirst,
608 stream_strategy: StreamStrategy::LeftThisEpoch,
609 edges: vec![JoinEdge {
610 left: JoinTable(1),
611 right: JoinTable(2),
612 left_join_key: vec![0],
613 right_join_key: vec![0],
614 }],
615 join_order: vec![JoinTable(1), JoinTable(2), JoinTable(3)],
616 };
617
618 let result = solver.solve();
619
620 assert!(result.is_err());
621 }
622}