polars_ops/frame/join/
cross_join.rs

1use polars_core::utils::{
2    _set_partition_size, CustomIterTools, NoNull, accumulate_dataframes_vertical_unchecked,
3    concat_df_unchecked, par_iter_bounded, split,
4};
5use polars_defs::join::{CrossJoinOptions, JoinType, MaintainOrderJoin};
6use polars_utils::pl_str::PlSmallStr;
7
8use super::*;
9
10fn slice_take(
11    total_rows: IdxSize,
12    n_rows_right: IdxSize,
13    slice: Option<(i64, usize)>,
14    inner: fn(IdxSize, IdxSize, IdxSize) -> IdxCa,
15) -> IdxCa {
16    match slice {
17        None => inner(0, total_rows, n_rows_right),
18        Some((offset, len)) => {
19            let (offset, len) = slice_offsets(offset, len, total_rows as usize);
20            inner(offset as IdxSize, (len + offset) as IdxSize, n_rows_right)
21        },
22    }
23}
24
25fn take_left(total_rows: IdxSize, n_rows_right: IdxSize, slice: Option<(i64, usize)>) -> IdxCa {
26    fn inner(offset: IdxSize, total_rows: IdxSize, n_rows_right: IdxSize) -> IdxCa {
27        let mut take: NoNull<IdxCa> = (offset..total_rows)
28            .map(|i| i / n_rows_right)
29            .collect_trusted();
30        take.set_sorted_flag(IsSorted::Ascending);
31        take.into_inner()
32    }
33    slice_take(total_rows, n_rows_right, slice, inner)
34}
35
36fn take_right(total_rows: IdxSize, n_rows_right: IdxSize, slice: Option<(i64, usize)>) -> IdxCa {
37    fn inner(offset: IdxSize, total_rows: IdxSize, n_rows_right: IdxSize) -> IdxCa {
38        let take: NoNull<IdxCa> = (offset..total_rows)
39            .map(|i| i % n_rows_right)
40            .collect_trusted();
41        take.into_inner()
42    }
43    slice_take(total_rows, n_rows_right, slice, inner)
44}
45
46pub trait CrossJoin: IntoDf {
47    /// Creates the Cartesian product from both frames, preserves the order of the left keys.
48    fn cross_join(
49        &self,
50        other: &DataFrame,
51        suffix: Option<PlSmallStr>,
52        slice: Option<(i64, usize)>,
53        maintain_order: MaintainOrderJoin,
54    ) -> PolarsResult<DataFrame> {
55        let (l_df, r_df) = cross_join_dfs(self.to_df(), other, slice, true, maintain_order)?;
56
57        _finish_join(l_df, r_df, suffix)
58    }
59}
60
61impl CrossJoin for DataFrame {}
62
63fn cross_join_dfs<'a>(
64    mut df_self: &'a DataFrame,
65    mut other: &'a DataFrame,
66    slice: Option<(i64, usize)>,
67    parallel: bool,
68    maintain_order: MaintainOrderJoin,
69) -> PolarsResult<(DataFrame, DataFrame)> {
70    if df_self.height() == 0 || other.height() == 0 {
71        return Ok((df_self.clear(), other.clear()));
72    }
73
74    let left_is_primary = match maintain_order {
75        MaintainOrderJoin::None => true,
76        MaintainOrderJoin::Left | MaintainOrderJoin::LeftRight => true,
77        MaintainOrderJoin::Right | MaintainOrderJoin::RightLeft => false,
78    };
79
80    if !left_is_primary {
81        core::mem::swap(&mut df_self, &mut other);
82    }
83
84    let n_rows_left = df_self.height() as IdxSize;
85    let n_rows_right = other.height() as IdxSize;
86    let Some(total_rows) = n_rows_left.checked_mul(n_rows_right) else {
87        polars_bail!(
88            ComputeError: "cross joins would produce more rows than fits into 2^32; \
89            consider compiling with polars-big-idx feature, or set 'streaming'"
90        );
91    };
92
93    // the left side has the Nth row combined with every row from right.
94    // So let's say we have the following no. of rows
95    // left: 3
96    // right: 4
97    //
98    // left take idx:   000011112222
99    // right take idx:  012301230123
100
101    let create_left_df = || {
102        // SAFETY:
103        // take left is in bounds
104        unsafe {
105            df_self.take_unchecked_impl(&take_left(total_rows, n_rows_right, slice), parallel)
106        }
107    };
108
109    let create_right_df = || {
110        // concatenation of dataframes is very expensive if we need to make the series mutable
111        // many times, these are atomic operations
112        // so we choose a different strategy at > 100 rows (arbitrarily small number)
113        if n_rows_left > 100 || slice.is_some() {
114            // SAFETY:
115            // take right is in bounds
116            unsafe {
117                other.take_unchecked_impl(&take_right(total_rows, n_rows_right, slice), parallel)
118            }
119        } else {
120            let iter = (0..n_rows_left).map(|_| other);
121            concat_df_unchecked(iter)
122        }
123    };
124    let (l_df, r_df) = if parallel {
125        try_raise_polars_abort();
126        RAYON.install(|| rayon::join(create_left_df, create_right_df))
127    } else {
128        (create_left_df(), create_right_df())
129    };
130    if left_is_primary {
131        Ok((l_df, r_df))
132    } else {
133        Ok((r_df, l_df))
134    }
135}
136
137pub(super) fn fused_cross_filter(
138    left: &DataFrame,
139    right: &DataFrame,
140    suffix: Option<PlSmallStr>,
141    cross_join_options: &CrossJoinOptions,
142    maintain_order: MaintainOrderJoin,
143    // If this is set, we didn't do a `how=cross`, but `how=left,anti,semi` in `join_where`
144    emit_unmatched_left: bool,
145    how: &JoinType,
146) -> PolarsResult<DataFrame> {
147    let unfiltered_size = (left.height() as u64).saturating_mul(right.height() as u64);
148    let chunk_size = (unfiltered_size / _set_partition_size() as u64).clamp(1, 100_000);
149    let num_chunks = (unfiltered_size / chunk_size).max(1) as usize;
150
151    let left_is_primary = match maintain_order {
152        MaintainOrderJoin::None => true,
153        MaintainOrderJoin::Left | MaintainOrderJoin::LeftRight => true,
154        MaintainOrderJoin::Right | MaintainOrderJoin::RightLeft => false,
155    };
156    // Every left row's full run of candidate matches must stay within a single chunk so a
157    // match can be decided per-chunk; that only holds when left is the (chunked) primary side.
158    polars_ensure!(
159        !emit_unmatched_left || left_is_primary,
160        InvalidOperation: "'maintain_order={:?}' is not supported for `join_where` with 'how'={}",
161        maintain_order,
162        how
163    );
164
165    let split_chunks;
166    let cartesian_prod = if left_is_primary {
167        split_chunks = split(left, num_chunks);
168        split_chunks.iter().map(|l| (l, right)).collect::<Vec<_>>()
169    } else {
170        split_chunks = split(right, num_chunks);
171        split_chunks.iter().map(|r| (left, r)).collect::<Vec<_>>()
172    };
173
174    let names = _finish_join(left.clear(), right.clear(), suffix)?;
175    let rename_names = names.get_column_names();
176    let rename_names = &rename_names[left.width()..];
177    let len_right = right.height();
178
179    let dfs = RAYON.install(|| {
180        par_iter_bounded(&cartesian_prod)
181            .map(|(left_chunk, right_chunk)| {
182                let (mut joined, right_taken) =
183                    cross_join_dfs(left_chunk, right_chunk, None, false, maintain_order)?;
184                let mut right_columns = right_taken.into_columns();
185
186                for (c, name) in right_columns.iter_mut().zip(rename_names) {
187                    c.rename((*name).clone());
188                }
189
190                unsafe { joined.hstack_mut_unchecked(&right_columns) };
191
192                if !emit_unmatched_left {
193                    cross_join_options.predicate.apply(joined, false)
194                } else {
195                    let mask = cross_join_options
196                        .predicate
197                        .evaluate(&joined)?
198                        .broadcast_owned_to(joined.height())?;
199
200                    let len_left = left_chunk.height();
201                    debug_assert_eq!(joined.height(), len_left * len_right);
202
203                    // Combine values and validity into one bitmap so a null bit reads as "no
204                    // match" (this is what filter has filtered)
205                    let mask_arr = mask.rechunk();
206                    let mask_arr = mask_arr.downcast_get(0).unwrap();
207                    let match_bits = match mask_arr.validity() {
208                        Some(validity) => mask_arr.values() & validity,
209                        None => mask_arr.values().clone(),
210                    };
211
212                    // Emit each left row's matches, or a single null-extended row when it has
213                    // none, so that unmatched rows keep their position in the left input's
214                    // order instead of being appended after the matched ones.
215                    let capacity = match_bits.set_bits() + len_left;
216                    let mut left_idx: Vec<IdxSize> = Vec::with_capacity(capacity);
217                    let mut right_idx: Vec<NullableIdxSize> = Vec::with_capacity(capacity);
218                    for i in 0..len_left {
219                        let run = match_bits.clone().sliced(i * len_right, len_right);
220                        if run.unset_bits() == len_right {
221                            left_idx.push(i as IdxSize);
222                            right_idx.push(NullableIdxSize::null());
223                        } else {
224                            for j in run.true_idx_iter() {
225                                left_idx.push(i as IdxSize);
226                                right_idx.push(NullableIdxSize::from(j as IdxSize));
227                            }
228                        }
229                    }
230
231                    let idx = unsafe { IdxCa::mmap_slice(PlSmallStr::EMPTY, &left_idx) };
232                    let mut out = unsafe { left_chunk.take_unchecked(&idx) };
233                    let mut right_columns = unsafe {
234                        IdxCa::with_nullable_idx(&right_idx, |idx| right.take_unchecked(idx))
235                    }
236                    .into_columns();
237                    for (c, name) in right_columns.iter_mut().zip(rename_names) {
238                        c.rename((*name).clone());
239                    }
240                    unsafe { out.hstack_mut_unchecked(&right_columns) };
241
242                    Ok(out)
243                }
244            })
245            .collect::<PolarsResult<Vec<_>>>()
246    })?;
247
248    Ok(accumulate_dataframes_vertical_unchecked(dfs))
249}
polars_ops/frame/join/cross_join.rs

polars_ops/frame/join/
cross_join.rs
