polars_ops/frame/join/hash_join/
mod.rs

1#![allow(unsafe_op_in_unsafe_fn)]
2pub(super) mod single_keys;
3mod single_keys_dispatch;
4mod single_keys_inner;
5mod single_keys_left;
6mod single_keys_outer;
7#[cfg(feature = "semi_anti_join")]
8mod single_keys_semi_anti;
9pub(super) mod sort_merge;
10use polars_arrow::array::ArrayRef;
11use polars_core::runtime::RAYON;
12use polars_core::utils::_set_partition_size;
13use polars_defs::join::{JoinArgs, JoinType, MaintainOrderJoin};
14use polars_utils::index::ChunkId;
15use polars_utils::unique_column_name;
16pub(super) use single_keys::*;
17pub use single_keys_dispatch::SeriesJoin;
18#[cfg(feature = "asof_join")]
19pub(super) use single_keys_dispatch::prepare_binary;
20use single_keys_inner::*;
21use single_keys_left::*;
22use single_keys_outer::*;
23#[cfg(feature = "semi_anti_join")]
24use single_keys_semi_anti::*;
25pub(crate) use sort_merge::*;
26
27pub use super::*;
28#[cfg(feature = "chunked_ids")]
29use crate::chunked_array::gather::chunked::TakeChunkedHorPar;
30
31pub fn default_join_ids() -> ChunkJoinOptIds {
32    #[cfg(feature = "chunked_ids")]
33    {
34        Either::Left(vec![])
35    }
36    #[cfg(not(feature = "chunked_ids"))]
37    {
38        vec![]
39    }
40}
41
42macro_rules! det_hash_prone_order {
43    ($self:expr, $other:expr) => {{
44        // The shortest relation will be used to create a hash table.
45        if $self.len() > $other.len() {
46            ($self, $other, false)
47        } else {
48            ($other, $self, true)
49        }
50    }};
51}
52
53pub(super) use det_hash_prone_order;
54#[cfg(feature = "performant")]
55use polars_arrow::legacy::conversion::primitive_to_vec;
56
57pub trait JoinDispatch: IntoDf {
58    /// # Safety
59    /// Join tuples must be in bounds
60    #[cfg(feature = "chunked_ids")]
61    unsafe fn create_left_df_chunked(
62        &self,
63        chunk_ids: &[ChunkId],
64        left_join: bool,
65        was_sliced: bool,
66    ) -> DataFrame {
67        let df_self = self.to_df();
68
69        let left_join_no_duplicate_matches =
70            left_join && !was_sliced && chunk_ids.len() == df_self.height();
71
72        if left_join_no_duplicate_matches {
73            df_self.clone()
74        } else {
75            // left join keys are in ascending order
76            let sorted = if left_join {
77                IsSorted::Ascending
78            } else {
79                IsSorted::Not
80            };
81            df_self._take_chunked_unchecked_hor_par(chunk_ids, sorted)
82        }
83    }
84
85    /// # Safety
86    /// Join tuples must be in bounds
87    unsafe fn _create_left_df_from_slice(
88        &self,
89        join_tuples: &[IdxSize],
90        left_join: bool,
91        was_sliced: bool,
92        sorted_tuple_idx: bool,
93    ) -> DataFrame {
94        let df_self = self.to_df();
95
96        let left_join_no_duplicate_matches =
97            sorted_tuple_idx && left_join && !was_sliced && join_tuples.len() == df_self.height();
98
99        if left_join_no_duplicate_matches {
100            df_self.clone()
101        } else {
102            let sorted = if sorted_tuple_idx {
103                IsSorted::Ascending
104            } else {
105                IsSorted::Not
106            };
107
108            df_self._take_unchecked_slice_sorted(join_tuples, true, sorted)
109        }
110    }
111
112    #[cfg(feature = "semi_anti_join")]
113    /// # Safety
114    /// `idx` must be in bounds
115    unsafe fn _finish_anti_semi_join(
116        &self,
117        mut idx: &[IdxSize],
118        slice: Option<(i64, usize)>,
119    ) -> DataFrame {
120        let ca_self = self.to_df();
121        if let Some((offset, len)) = slice {
122            idx = slice_slice(idx, offset, len);
123        }
124        // idx from anti-semi join should always be sorted
125        ca_self._take_unchecked_slice_sorted(idx, true, IsSorted::Ascending)
126    }
127
128    #[cfg(feature = "semi_anti_join")]
129    fn _semi_anti_join_from_series(
130        &self,
131        s_left: &Series,
132        s_right: &Series,
133        slice: Option<(i64, usize)>,
134        anti: bool,
135        nulls_equal: bool,
136    ) -> PolarsResult<DataFrame> {
137        let ca_self = self.to_df();
138
139        let idx = s_left.hash_join_semi_anti(s_right, anti, nulls_equal)?;
140        // SAFETY:
141        // indices are in bounds
142        Ok(unsafe { ca_self._finish_anti_semi_join(&idx, slice) })
143    }
144    fn _full_join_from_series(
145        &self,
146        other: &DataFrame,
147        s_left: &Series,
148        s_right: &Series,
149        args: JoinArgs,
150    ) -> PolarsResult<DataFrame> {
151        let df_self = self.to_df();
152
153        // Get the indexes of the joined relations
154        let (mut join_idx_l, mut join_idx_r) =
155            s_left.hash_join_outer(s_right, args.validation, args.nulls_equal)?;
156
157        try_raise_polars_abort();
158
159        let (df_left, df_right) = if args.maintain_order != MaintainOrderJoin::None {
160            let idx_ca_l = IdxCa::with_chunk("a".into(), join_idx_l);
161            let idx_ca_r = IdxCa::with_chunk("b".into(), join_idx_r);
162            let mut df = unsafe {
163                DataFrame::new_unchecked_infer_height(vec![
164                    idx_ca_l.into_series().into(),
165                    idx_ca_r.into_series().into(),
166                ])
167            };
168
169            let options = SortMultipleOptions::new()
170                .with_order_descending(false)
171                .with_maintain_order(true)
172                .with_nulls_last(true);
173
174            let columns = match args.maintain_order {
175                MaintainOrderJoin::Left => vec!["a"],
176                MaintainOrderJoin::LeftRight => vec!["a", "b"],
177                MaintainOrderJoin::Right => vec!["b"],
178                MaintainOrderJoin::RightLeft => vec!["b", "a"],
179                _ => unreachable!(),
180            };
181
182            df.sort_in_place(columns, options)?;
183
184            // If the order is maintained, we can only slice after sorting
185            if let Some((offset, len)) = args.slice {
186                df = df.slice(offset, len);
187            }
188
189            let join_tuples_left = df.column("a").unwrap().idx().unwrap();
190            let join_tuples_right = df.column("b").unwrap().idx().unwrap();
191            RAYON.join(
192                || unsafe { df_self.take_unchecked(join_tuples_left) },
193                || unsafe { other.take_unchecked(join_tuples_right) },
194            )
195        } else {
196            if let Some((offset, len)) = args.slice {
197                let (offset, len) = slice_offsets(offset, len, join_idx_l.len());
198                join_idx_l.slice(offset, len);
199                join_idx_r.slice(offset, len);
200            }
201            let idx_ca_l = IdxCa::with_chunk("a".into(), join_idx_l);
202            let idx_ca_r = IdxCa::with_chunk("b".into(), join_idx_r);
203            RAYON.join(
204                || unsafe { df_self.take_unchecked(&idx_ca_l) },
205                || unsafe { other.take_unchecked(&idx_ca_r) },
206            )
207        };
208
209        let coalesce = args.coalesce.coalesce(&JoinType::Full);
210        if coalesce {
211            let tmp_right_name = unique_column_name();
212            let mut df_right = df_right;
213            df_right.rename(s_right.name().as_str(), tmp_right_name.clone())?;
214            let out = _finish_join(df_left, df_right, args.suffix.clone())?;
215            Ok(_coalesce_full_join(
216                out,
217                &[s_left.name().clone()],
218                &[tmp_right_name],
219                args.suffix,
220                df_self,
221            ))
222        } else {
223            _finish_join(df_left, df_right, args.suffix.clone())
224        }
225    }
226}
227
228impl JoinDispatch for DataFrame {}
polars_ops/frame/join/hash_join/mod.rs

polars_ops/frame/join/hash_join/
mod.rs
