Skip to main content

oxirs_arq/streaming/
streaminghashjoin_traits.rs

1//! # StreamingHashJoin - Trait Implementations
2//!
3//! This module contains trait implementations for `StreamingHashJoin`.
4//!
5//! ## Implemented Traits
6//!
7//! - `DataStream`
8//!
9//! 🤖 Generated with [SplitRS](https://github.com/cool-japan/splitrs)
10
11use crate::algebra::Solution;
12use anyhow::{anyhow, Result};
13
14use super::functions::DataStream;
15use super::types::{StreamStats, StreamingHashJoin};
16
17impl DataStream for StreamingHashJoin {
18    fn next_batch(&mut self) -> Result<Option<Vec<Solution>>> {
19        if !self.left_exhausted {
20            while let Some(batch) = self.left_stream.next_batch()? {
21                for solution in batch {
22                    let key = self.extract_join_key(&solution);
23                    let estimated_size = std::mem::size_of_val(&solution) + key.len();
24                    if !self
25                        .memory_monitor
26                        .allocate(estimated_size, "hash_join_build")
27                    {
28                        self.spill_hash_table()?;
29                        if !self
30                            .memory_monitor
31                            .allocate(estimated_size, "hash_join_build")
32                        {
33                            return Err(anyhow!("Cannot allocate memory even after spilling"));
34                        }
35                    }
36                    self.hash_table.entry(key).or_default().push(solution);
37                }
38            }
39            self.left_exhausted = true;
40        }
41        if let Some(right_batch) = self.right_stream.next_batch()? {
42            let mut result_batch = Vec::new();
43            for right_solution in right_batch {
44                let key = self.extract_join_key(&right_solution);
45                if let Some(left_solutions) = self.hash_table.get(&key) {
46                    for left_solution in left_solutions {
47                        if let Some(joined) = self.join_solutions(left_solution, &right_solution) {
48                            result_batch.push(joined);
49                        }
50                    }
51                }
52            }
53            Ok(if result_batch.is_empty() {
54                None
55            } else {
56                Some(result_batch)
57            })
58        } else {
59            Ok(None)
60        }
61    }
62    fn has_more(&self) -> bool {
63        !self.left_exhausted || self.right_stream.has_more()
64    }
65    fn estimated_size(&self) -> Option<usize> {
66        None
67    }
68    fn reset(&mut self) -> Result<()> {
69        self.left_stream.reset()?;
70        self.right_stream.reset()?;
71        self.hash_table.clear();
72        self.left_exhausted = false;
73        self.current_batch = None;
74        Ok(())
75    }
76    fn get_stats(&self) -> StreamStats {
77        StreamStats::default()
78    }
79}