polars_ops/frame/join/hash_join/
single_keys_dispatch.rs

1use polars_arrow::array::PrimitiveArray;
2use polars_core::chunked_array::ops::row_encode::{
3    encode_rows_unordered, encode_rows_vertical_par_unordered_broadcast_nulls,
4};
5use polars_core::series::BitRepr;
6use polars_core::utils::split;
7use polars_core::with_match_physical_float_polars_type;
8use polars_defs::join::JoinValidation;
9use polars_utils::aliases::PlRandomState;
10use polars_utils::hashing::DirtyHash;
11use polars_utils::nulls::IsNull;
12use polars_utils::total_ord::{ToTotalOrd, TotalEq, TotalHash};
13
14use super::*;
15use crate::frame::join::validation::validate_probe;
16use crate::series::SeriesSealed;
17
18pub trait SeriesJoin: SeriesSealed + Sized {
19    #[doc(hidden)]
20    fn hash_join_left(
21        &self,
22        other: &Series,
23        validate: JoinValidation,
24        nulls_equal: bool,
25    ) -> PolarsResult<LeftJoinIds> {
26        let s_self = self.as_series();
27        let (lhs, rhs) = (s_self.to_physical_repr(), other.to_physical_repr());
28        validate_probe(validate, &lhs, &rhs, false, nulls_equal)?;
29
30        let lhs_dtype = lhs.dtype();
31        let rhs_dtype = rhs.dtype();
32
33        use DataType as T;
34        match lhs_dtype {
35            T::String | T::Binary => {
36                let lhs = lhs.cast(&T::Binary).unwrap();
37                let rhs = rhs.cast(&T::Binary).unwrap();
38                let lhs = lhs.binary().unwrap();
39                let rhs = rhs.binary().unwrap();
40                let (lhs, rhs, _, _) = prepare_binary::<BinaryType>(lhs, rhs, false);
41                let lhs = lhs.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
42                let rhs = rhs.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
43                let build_null_count = other.null_count();
44                hash_join_tuples_left(
45                    lhs,
46                    rhs,
47                    None,
48                    None,
49                    validate,
50                    nulls_equal,
51                    build_null_count,
52                )
53            },
54            T::BinaryOffset => {
55                let lhs = lhs.binary_offset().unwrap();
56                let rhs = rhs.binary_offset().unwrap();
57                let (lhs, rhs, _, _) = prepare_binary::<BinaryOffsetType>(lhs, rhs, false);
58                // Take slices so that vecs are not copied
59                let lhs = lhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
60                let rhs = rhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
61                let build_null_count = other.null_count();
62                hash_join_tuples_left(
63                    lhs,
64                    rhs,
65                    None,
66                    None,
67                    validate,
68                    nulls_equal,
69                    build_null_count,
70                )
71            },
72            T::List(_) => {
73                let lhs = &encode_join_nested_key(lhs.into_owned(), nulls_equal)?;
74                let rhs = &encode_join_nested_key(rhs.into_owned(), nulls_equal)?;
75                lhs.hash_join_left(rhs, validate, nulls_equal)
76            },
77            #[cfg(feature = "dtype-array")]
78            T::Array(_, _) => {
79                let lhs = &encode_join_nested_key(lhs.into_owned(), nulls_equal)?;
80                let rhs = &encode_join_nested_key(rhs.into_owned(), nulls_equal)?;
81                lhs.hash_join_left(rhs, validate, nulls_equal)
82            },
83            #[cfg(feature = "dtype-struct")]
84            T::Struct(_) => {
85                let lhs = &encode_join_nested_key(lhs.into_owned(), nulls_equal)?;
86                let rhs = &encode_join_nested_key(rhs.into_owned(), nulls_equal)?;
87                lhs.hash_join_left(rhs, validate, nulls_equal)
88            },
89            x if x.is_float() => {
90                with_match_physical_float_polars_type!(lhs.dtype(), |$T| {
91                    let lhs: &ChunkedArray<$T> = lhs.as_ref().as_ref().as_ref();
92                    let rhs: &ChunkedArray<$T> = rhs.as_ref().as_ref().as_ref();
93                    num_group_join_left(lhs, rhs, validate, nulls_equal)
94                })
95            },
96            _ => {
97                let lhs = s_self.bit_repr();
98                let rhs = other.bit_repr();
99
100                let (Some(lhs), Some(rhs)) = (lhs, rhs) else {
101                    polars_bail!(nyi = "Hash Left Join between {lhs_dtype} and {rhs_dtype}");
102                };
103
104                use BitRepr as B;
105                match (lhs, rhs) {
106                    (B::U8(lhs), B::U8(rhs)) => {
107                        num_group_join_left(&lhs, &rhs, validate, nulls_equal)
108                    },
109                    (B::U16(lhs), B::U16(rhs)) => {
110                        num_group_join_left(&lhs, &rhs, validate, nulls_equal)
111                    },
112                    (B::U32(lhs), B::U32(rhs)) => {
113                        num_group_join_left(&lhs, &rhs, validate, nulls_equal)
114                    },
115                    (B::U64(lhs), B::U64(rhs)) => {
116                        num_group_join_left(&lhs, &rhs, validate, nulls_equal)
117                    },
118                    #[cfg(feature = "dtype-u128")]
119                    (B::U128(lhs), B::U128(rhs)) => {
120                        num_group_join_left(&lhs, &rhs, validate, nulls_equal)
121                    },
122                    _ => {
123                        polars_bail!(
124                            nyi = "Mismatch bit repr Hash Left Join between {lhs_dtype} and {rhs_dtype}",
125                        );
126                    },
127                }
128            },
129        }
130    }
131
132    #[cfg(feature = "semi_anti_join")]
133    fn hash_join_semi_anti(
134        &self,
135        other: &Series,
136        anti: bool,
137        nulls_equal: bool,
138    ) -> PolarsResult<Vec<IdxSize>> {
139        let s_self = self.as_series();
140        let (lhs, rhs) = (s_self.to_physical_repr(), other.to_physical_repr());
141
142        let lhs_dtype = lhs.dtype();
143        let rhs_dtype = rhs.dtype();
144
145        use DataType as T;
146        Ok(match lhs_dtype {
147            T::String | T::Binary => {
148                let lhs = lhs.cast(&T::Binary).unwrap();
149                let rhs = rhs.cast(&T::Binary).unwrap();
150                let lhs = lhs.binary().unwrap();
151                let rhs = rhs.binary().unwrap();
152                let (lhs, rhs, _, _) = prepare_binary::<BinaryType>(lhs, rhs, false);
153                // Take slices so that vecs are not copied
154                let lhs = lhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
155                let rhs = rhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
156                if anti {
157                    hash_join_tuples_left_anti(lhs, rhs, nulls_equal)
158                } else {
159                    hash_join_tuples_left_semi(lhs, rhs, nulls_equal)
160                }
161            },
162            T::BinaryOffset => {
163                let lhs = lhs.binary_offset().unwrap();
164                let rhs = rhs.binary_offset().unwrap();
165                let (lhs, rhs, _, _) = prepare_binary::<BinaryOffsetType>(lhs, rhs, false);
166                // Take slices so that vecs are not copied
167                let lhs = lhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
168                let rhs = rhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
169                if anti {
170                    hash_join_tuples_left_anti(lhs, rhs, nulls_equal)
171                } else {
172                    hash_join_tuples_left_semi(lhs, rhs, nulls_equal)
173                }
174            },
175            T::List(_) => {
176                let lhs = &encode_join_nested_key(lhs.into_owned(), nulls_equal)?;
177                let rhs = &encode_join_nested_key(rhs.into_owned(), nulls_equal)?;
178                lhs.hash_join_semi_anti(rhs, anti, nulls_equal)?
179            },
180            #[cfg(feature = "dtype-array")]
181            T::Array(_, _) => {
182                let lhs = &encode_join_nested_key(lhs.into_owned(), nulls_equal)?;
183                let rhs = &encode_join_nested_key(rhs.into_owned(), nulls_equal)?;
184                lhs.hash_join_semi_anti(rhs, anti, nulls_equal)?
185            },
186            #[cfg(feature = "dtype-struct")]
187            T::Struct(_) => {
188                let lhs = &encode_join_nested_key(lhs.into_owned(), nulls_equal)?;
189                let rhs = &encode_join_nested_key(rhs.into_owned(), nulls_equal)?;
190                lhs.hash_join_semi_anti(rhs, anti, nulls_equal)?
191            },
192            x if x.is_float() => {
193                with_match_physical_float_polars_type!(lhs.dtype(), |$T| {
194                    let lhs: &ChunkedArray<$T> = lhs.as_ref().as_ref().as_ref();
195                    let rhs: &ChunkedArray<$T> = rhs.as_ref().as_ref().as_ref();
196                    num_group_join_anti_semi(lhs, rhs, anti, nulls_equal)
197                })
198            },
199            _ => {
200                let lhs = s_self.bit_repr();
201                let rhs = other.bit_repr();
202
203                let (Some(lhs), Some(rhs)) = (lhs, rhs) else {
204                    polars_bail!(nyi = "Hash Semi-Anti Join between {lhs_dtype} and {rhs_dtype}");
205                };
206
207                use BitRepr as B;
208                match (lhs, rhs) {
209                    (B::U8(lhs), B::U8(rhs)) => {
210                        num_group_join_anti_semi(&lhs, &rhs, anti, nulls_equal)
211                    },
212                    (B::U16(lhs), B::U16(rhs)) => {
213                        num_group_join_anti_semi(&lhs, &rhs, anti, nulls_equal)
214                    },
215                    (B::U32(lhs), B::U32(rhs)) => {
216                        num_group_join_anti_semi(&lhs, &rhs, anti, nulls_equal)
217                    },
218                    (B::U64(lhs), B::U64(rhs)) => {
219                        num_group_join_anti_semi(&lhs, &rhs, anti, nulls_equal)
220                    },
221                    #[cfg(feature = "dtype-u128")]
222                    (B::U128(lhs), B::U128(rhs)) => {
223                        num_group_join_anti_semi(&lhs, &rhs, anti, nulls_equal)
224                    },
225                    _ => {
226                        polars_bail!(
227                            nyi = "Mismatch bit repr Hash Semi-Anti Join between {lhs_dtype} and {rhs_dtype}",
228                        );
229                    },
230                }
231            },
232        })
233    }
234
235    // returns the join tuples and whether or not the lhs tuples are sorted
236    fn hash_join_inner(
237        &self,
238        other: &Series,
239        validate: JoinValidation,
240        nulls_equal: bool,
241    ) -> PolarsResult<(InnerJoinIds, bool)> {
242        let s_self = self.as_series();
243        let (lhs, rhs) = (s_self.to_physical_repr(), other.to_physical_repr());
244        validate_probe(validate, &lhs, &rhs, true, nulls_equal)?;
245
246        let lhs_dtype = lhs.dtype();
247        let rhs_dtype = rhs.dtype();
248
249        use DataType as T;
250        match lhs_dtype {
251            T::String | T::Binary => {
252                let lhs = lhs.cast(&T::Binary).unwrap();
253                let rhs = rhs.cast(&T::Binary).unwrap();
254                let lhs = lhs.binary().unwrap();
255                let rhs = rhs.binary().unwrap();
256                let (lhs, rhs, swapped, _) = prepare_binary::<BinaryType>(lhs, rhs, true);
257                // Take slices so that vecs are not copied
258                let lhs = lhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
259                let rhs = rhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
260                let build_null_count = if swapped {
261                    s_self.null_count()
262                } else {
263                    other.null_count()
264                };
265                Ok((
266                    hash_join_tuples_inner(
267                        lhs,
268                        rhs,
269                        swapped,
270                        validate,
271                        nulls_equal,
272                        build_null_count,
273                    )?,
274                    !swapped,
275                ))
276            },
277            T::BinaryOffset => {
278                let lhs = lhs.binary_offset().unwrap();
279                let rhs = rhs.binary_offset()?;
280                let (lhs, rhs, swapped, _) = prepare_binary::<BinaryOffsetType>(lhs, rhs, true);
281                // Take slices so that vecs are not copied
282                let lhs = lhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
283                let rhs = rhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
284                let build_null_count = if swapped {
285                    s_self.null_count()
286                } else {
287                    other.null_count()
288                };
289                Ok((
290                    hash_join_tuples_inner(
291                        lhs,
292                        rhs,
293                        swapped,
294                        validate,
295                        nulls_equal,
296                        build_null_count,
297                    )?,
298                    !swapped,
299                ))
300            },
301            T::List(_) => {
302                let lhs = &encode_join_nested_key(lhs.into_owned(), nulls_equal)?;
303                let rhs = &encode_join_nested_key(rhs.into_owned(), nulls_equal)?;
304                lhs.hash_join_inner(rhs, validate, nulls_equal)
305            },
306            #[cfg(feature = "dtype-array")]
307            T::Array(_, _) => {
308                let lhs = &encode_join_nested_key(lhs.into_owned(), nulls_equal)?;
309                let rhs = &encode_join_nested_key(rhs.into_owned(), nulls_equal)?;
310                lhs.hash_join_inner(rhs, validate, nulls_equal)
311            },
312            #[cfg(feature = "dtype-struct")]
313            T::Struct(_) => {
314                let lhs = &encode_join_nested_key(lhs.into_owned(), nulls_equal)?;
315                let rhs = &encode_join_nested_key(rhs.into_owned(), nulls_equal)?;
316                lhs.hash_join_inner(rhs, validate, nulls_equal)
317            },
318            x if x.is_float() => {
319                with_match_physical_float_polars_type!(lhs.dtype(), |$T| {
320                    let lhs: &ChunkedArray<$T> = lhs.as_ref().as_ref().as_ref();
321                    let rhs: &ChunkedArray<$T> = rhs.as_ref().as_ref().as_ref();
322                    group_join_inner::<$T>(lhs, rhs, validate, nulls_equal)
323                })
324            },
325            _ => {
326                let lhs = s_self.bit_repr();
327                let rhs = other.bit_repr();
328
329                let (Some(lhs), Some(rhs)) = (lhs, rhs) else {
330                    polars_bail!(nyi = "Hash Inner Join between {lhs_dtype} and {rhs_dtype}");
331                };
332
333                use BitRepr as B;
334                match (lhs, rhs) {
335                    (B::U8(lhs), B::U8(rhs)) => {
336                        group_join_inner::<UInt8Type>(&lhs, &rhs, validate, nulls_equal)
337                    },
338                    (B::U16(lhs), B::U16(rhs)) => {
339                        group_join_inner::<UInt16Type>(&lhs, &rhs, validate, nulls_equal)
340                    },
341                    (B::U32(lhs), B::U32(rhs)) => {
342                        group_join_inner::<UInt32Type>(&lhs, &rhs, validate, nulls_equal)
343                    },
344                    (B::U64(lhs), BitRepr::U64(rhs)) => {
345                        group_join_inner::<UInt64Type>(&lhs, &rhs, validate, nulls_equal)
346                    },
347                    #[cfg(feature = "dtype-u128")]
348                    (B::U128(lhs), BitRepr::U128(rhs)) => {
349                        group_join_inner::<UInt128Type>(&lhs, &rhs, validate, nulls_equal)
350                    },
351                    _ => {
352                        polars_bail!(
353                            nyi = "Mismatch bit repr Hash Inner Join between {lhs_dtype} and {rhs_dtype}"
354                        );
355                    },
356                }
357            },
358        }
359    }
360
361    fn hash_join_outer(
362        &self,
363        other: &Series,
364        validate: JoinValidation,
365        nulls_equal: bool,
366    ) -> PolarsResult<(PrimitiveArray<IdxSize>, PrimitiveArray<IdxSize>)> {
367        let s_self = self.as_series();
368        let (lhs, rhs) = (s_self.to_physical_repr(), other.to_physical_repr());
369        validate_probe(validate, &lhs, &rhs, true, nulls_equal)?;
370
371        let lhs_dtype = lhs.dtype();
372        let rhs_dtype = rhs.dtype();
373
374        use DataType as T;
375        match lhs_dtype {
376            T::String | T::Binary => {
377                let lhs = lhs.cast(&T::Binary).unwrap();
378                let rhs = rhs.cast(&T::Binary).unwrap();
379                let lhs = lhs.binary().unwrap();
380                let rhs = rhs.binary().unwrap();
381                let (lhs, rhs, swapped, _) = prepare_binary::<BinaryType>(lhs, rhs, true);
382                // Take slices so that vecs are not copied
383                let lhs = lhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
384                let rhs = rhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
385                hash_join_tuples_outer(lhs, rhs, swapped, validate, nulls_equal)
386            },
387            T::BinaryOffset => {
388                let lhs = lhs.binary_offset().unwrap();
389                let rhs = rhs.binary_offset()?;
390                let (lhs, rhs, swapped, _) = prepare_binary::<BinaryOffsetType>(lhs, rhs, true);
391                // Take slices so that vecs are not copied
392                let lhs = lhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
393                let rhs = rhs.iter().map(|k| k.as_slice()).collect::<Vec<_>>();
394                hash_join_tuples_outer(lhs, rhs, swapped, validate, nulls_equal)
395            },
396            T::List(_) => {
397                let lhs = &encode_join_nested_key(lhs.into_owned(), nulls_equal)?;
398                let rhs = &encode_join_nested_key(rhs.into_owned(), nulls_equal)?;
399                lhs.hash_join_outer(rhs, validate, nulls_equal)
400            },
401            #[cfg(feature = "dtype-array")]
402            T::Array(_, _) => {
403                let lhs = &encode_join_nested_key(lhs.into_owned(), nulls_equal)?;
404                let rhs = &encode_join_nested_key(rhs.into_owned(), nulls_equal)?;
405                lhs.hash_join_outer(rhs, validate, nulls_equal)
406            },
407            #[cfg(feature = "dtype-struct")]
408            T::Struct(_) => {
409                let lhs = &encode_join_nested_key(lhs.into_owned(), nulls_equal)?;
410                let rhs = &encode_join_nested_key(rhs.into_owned(), nulls_equal)?;
411                lhs.hash_join_outer(rhs, validate, nulls_equal)
412            },
413            x if x.is_float() => {
414                with_match_physical_float_polars_type!(lhs.dtype(), |$T| {
415                    let lhs: &ChunkedArray<$T> = lhs.as_ref().as_ref().as_ref();
416                    let rhs: &ChunkedArray<$T> = rhs.as_ref().as_ref().as_ref();
417                    hash_join_outer(lhs, rhs, validate, nulls_equal)
418                })
419            },
420            _ => {
421                let (Some(lhs), Some(rhs)) = (s_self.bit_repr(), other.bit_repr()) else {
422                    polars_bail!(nyi = "Hash Join Outer between {lhs_dtype} and {rhs_dtype}");
423                };
424
425                use BitRepr as B;
426                match (lhs, rhs) {
427                    (B::U8(lhs), B::U8(rhs)) => hash_join_outer(&lhs, &rhs, validate, nulls_equal),
428                    (B::U16(lhs), B::U16(rhs)) => {
429                        hash_join_outer(&lhs, &rhs, validate, nulls_equal)
430                    },
431                    (B::U32(lhs), B::U32(rhs)) => {
432                        hash_join_outer(&lhs, &rhs, validate, nulls_equal)
433                    },
434                    (B::U64(lhs), B::U64(rhs)) => {
435                        hash_join_outer(&lhs, &rhs, validate, nulls_equal)
436                    },
437                    #[cfg(feature = "dtype-u128")]
438                    (B::U128(lhs), B::U128(rhs)) => {
439                        hash_join_outer(&lhs, &rhs, validate, nulls_equal)
440                    },
441                    _ => {
442                        polars_bail!(
443                            nyi = "Mismatch bit repr Hash Join Outer between {lhs_dtype} and {rhs_dtype}"
444                        );
445                    },
446                }
447            },
448        }
449    }
450}
451
452impl SeriesJoin for Series {}
453
454fn chunks_as_slices<T>(splitted: &[ChunkedArray<T>]) -> Vec<&[T::Native]>
455where
456    T: PolarsNumericType,
457{
458    splitted
459        .iter()
460        .flat_map(|ca| ca.downcast_iter().map(|arr| arr.values().as_slice()))
461        .collect()
462}
463
464fn encode_join_nested_key(s: Series, nulls_equal: bool) -> PolarsResult<Series> {
465    let by = [s.into_column()];
466    let encoded = if nulls_equal {
467        encode_rows_unordered(&by)?
468    } else {
469        encode_rows_vertical_par_unordered_broadcast_nulls(&by)?
470    };
471    Ok(encoded.into_series())
472}
473
474fn get_arrays<T: PolarsDataType>(cas: &[ChunkedArray<T>]) -> Vec<&T::Array> {
475    cas.iter().flat_map(|arr| arr.downcast_iter()).collect()
476}
477
478fn group_join_inner<T>(
479    left: &ChunkedArray<T>,
480    right: &ChunkedArray<T>,
481    validate: JoinValidation,
482    nulls_equal: bool,
483) -> PolarsResult<(InnerJoinIds, bool)>
484where
485    T: PolarsDataType,
486    for<'a> &'a T::Array: IntoIterator<Item = Option<&'a T::Physical<'a>>>,
487    for<'a> T::Physical<'a>:
488        Send + Sync + Copy + TotalHash + TotalEq + DirtyHash + IsNull + ToTotalOrd,
489    for<'a> <T::Physical<'a> as ToTotalOrd>::TotalOrdItem:
490        Send + Sync + Copy + Hash + Eq + DirtyHash + IsNull,
491{
492    let n_threads = RAYON.current_num_threads();
493    let (a, b, swapped) = det_hash_prone_order!(left, right);
494    let splitted_a = split(a, n_threads);
495    let splitted_b = split(b, n_threads);
496    let splitted_a = get_arrays(&splitted_a);
497    let splitted_b = get_arrays(&splitted_b);
498
499    match (left.null_count(), right.null_count()) {
500        (0, 0) => {
501            let first = &splitted_a[0];
502            if first.as_slice().is_some() {
503                let splitted_a = splitted_a
504                    .iter()
505                    .map(|arr| arr.as_slice().unwrap())
506                    .collect::<Vec<_>>();
507                let splitted_b = splitted_b
508                    .iter()
509                    .map(|arr| arr.as_slice().unwrap())
510                    .collect::<Vec<_>>();
511                Ok((
512                    hash_join_tuples_inner(
513                        splitted_a,
514                        splitted_b,
515                        swapped,
516                        validate,
517                        nulls_equal,
518                        0,
519                    )?,
520                    !swapped,
521                ))
522            } else {
523                Ok((
524                    hash_join_tuples_inner(
525                        splitted_a,
526                        splitted_b,
527                        swapped,
528                        validate,
529                        nulls_equal,
530                        0,
531                    )?,
532                    !swapped,
533                ))
534            }
535        },
536        _ => {
537            let build_null_count = if swapped {
538                left.null_count()
539            } else {
540                right.null_count()
541            };
542            Ok((
543                hash_join_tuples_inner(
544                    splitted_a,
545                    splitted_b,
546                    swapped,
547                    validate,
548                    nulls_equal,
549                    build_null_count,
550                )?,
551                !swapped,
552            ))
553        },
554    }
555}
556
557#[cfg(feature = "chunked_ids")]
558fn create_mappings(
559    chunks_left: &[ArrayRef],
560    chunks_right: &[ArrayRef],
561    left_len: usize,
562    right_len: usize,
563) -> (Option<Vec<ChunkId>>, Option<Vec<ChunkId>>) {
564    let mapping_left = || {
565        if chunks_left.len() > 1 {
566            Some(create_chunked_index_mapping(chunks_left, left_len))
567        } else {
568            None
569        }
570    };
571
572    let mapping_right = || {
573        if chunks_right.len() > 1 {
574            Some(create_chunked_index_mapping(chunks_right, right_len))
575        } else {
576            None
577        }
578    };
579
580    RAYON.join(mapping_left, mapping_right)
581}
582
583#[cfg(not(feature = "chunked_ids"))]
584fn create_mappings(
585    _chunks_left: &[ArrayRef],
586    _chunks_right: &[ArrayRef],
587    _left_len: usize,
588    _right_len: usize,
589) -> (Option<Vec<ChunkId>>, Option<Vec<ChunkId>>) {
590    (None, None)
591}
592
593fn num_group_join_left<T>(
594    left: &ChunkedArray<T>,
595    right: &ChunkedArray<T>,
596    validate: JoinValidation,
597    nulls_equal: bool,
598) -> PolarsResult<LeftJoinIds>
599where
600    T: PolarsNumericType,
601    T::Native: TotalHash + TotalEq + DirtyHash + IsNull + ToTotalOrd,
602    <T::Native as ToTotalOrd>::TotalOrdItem: Send + Sync + Copy + Hash + Eq + DirtyHash + IsNull,
603    T::Native: DirtyHash + Copy + ToTotalOrd,
604    <Option<T::Native> as ToTotalOrd>::TotalOrdItem: Send + Sync + DirtyHash,
605{
606    let n_threads = RAYON.current_num_threads();
607    let splitted_a = split(left, n_threads);
608    let splitted_b = split(right, n_threads);
609    match (
610        left.null_count(),
611        right.null_count(),
612        left.chunks().len(),
613        right.chunks().len(),
614    ) {
615        (0, 0, 1, 1) => {
616            let keys_a = chunks_as_slices(&splitted_a);
617            let keys_b = chunks_as_slices(&splitted_b);
618            hash_join_tuples_left(keys_a, keys_b, None, None, validate, nulls_equal, 0)
619        },
620        (0, 0, _, _) => {
621            let keys_a = chunks_as_slices(&splitted_a);
622            let keys_b = chunks_as_slices(&splitted_b);
623
624            let (mapping_left, mapping_right) =
625                create_mappings(left.chunks(), right.chunks(), left.len(), right.len());
626            hash_join_tuples_left(
627                keys_a,
628                keys_b,
629                mapping_left.as_deref(),
630                mapping_right.as_deref(),
631                validate,
632                nulls_equal,
633                0,
634            )
635        },
636        _ => {
637            let keys_a = get_arrays(&splitted_a);
638            let keys_b = get_arrays(&splitted_b);
639            let (mapping_left, mapping_right) =
640                create_mappings(left.chunks(), right.chunks(), left.len(), right.len());
641            let build_null_count = right.null_count();
642            hash_join_tuples_left(
643                keys_a,
644                keys_b,
645                mapping_left.as_deref(),
646                mapping_right.as_deref(),
647                validate,
648                nulls_equal,
649                build_null_count,
650            )
651        },
652    }
653}
654
655fn hash_join_outer<T>(
656    ca_in: &ChunkedArray<T>,
657    other: &ChunkedArray<T>,
658    validate: JoinValidation,
659    nulls_equal: bool,
660) -> PolarsResult<(PrimitiveArray<IdxSize>, PrimitiveArray<IdxSize>)>
661where
662    T: PolarsNumericType,
663    T::Native: TotalHash + TotalEq + ToTotalOrd,
664    <T::Native as ToTotalOrd>::TotalOrdItem: Send + Sync + Copy + Hash + Eq + IsNull,
665{
666    let (a, b, swapped) = det_hash_prone_order!(ca_in, other);
667
668    let n_partitions = _set_partition_size();
669    let splitted_a = split(a, n_partitions);
670    let splitted_b = split(b, n_partitions);
671
672    match (a.null_count(), b.null_count()) {
673        (0, 0) => {
674            let iters_a = splitted_a
675                .iter()
676                .flat_map(|ca| ca.downcast_iter().map(|arr| arr.values().as_slice()))
677                .collect::<Vec<_>>();
678            let iters_b = splitted_b
679                .iter()
680                .flat_map(|ca| ca.downcast_iter().map(|arr| arr.values().as_slice()))
681                .collect::<Vec<_>>();
682            hash_join_tuples_outer(iters_a, iters_b, swapped, validate, nulls_equal)
683        },
684        _ => {
685            let iters_a = splitted_a
686                .iter()
687                .flat_map(|ca| ca.downcast_iter().map(|arr| arr.iter()))
688                .collect::<Vec<_>>();
689            let iters_b = splitted_b
690                .iter()
691                .flat_map(|ca| ca.downcast_iter().map(|arr| arr.iter()))
692                .collect::<Vec<_>>();
693            hash_join_tuples_outer(iters_a, iters_b, swapped, validate, nulls_equal)
694        },
695    }
696}
697
698pub(crate) fn prepare_binary<'a, T>(
699    ca: &'a ChunkedArray<T>,
700    other: &'a ChunkedArray<T>,
701    // In inner join and outer join, the shortest relation will be used to create a hash table.
702    // In left join, always use the right side to create.
703    build_shortest_table: bool,
704) -> (
705    Vec<Vec<BytesHash<'a>>>,
706    Vec<Vec<BytesHash<'a>>>,
707    bool,
708    PlRandomState,
709)
710where
711    T: PolarsDataType,
712    for<'b> <T::Array as StaticArray>::ValueT<'b>: AsRef<[u8]>,
713{
714    let (a, b, swapped) = if build_shortest_table {
715        det_hash_prone_order!(ca, other)
716    } else {
717        (ca, other, false)
718    };
719    let hb = PlRandomState::default();
720    let bh_a = a.to_bytes_hashes(true, hb.clone());
721    let bh_b = b.to_bytes_hashes(true, hb.clone());
722
723    (bh_a, bh_b, swapped, hb)
724}
725
726#[cfg(feature = "semi_anti_join")]
727fn num_group_join_anti_semi<T>(
728    left: &ChunkedArray<T>,
729    right: &ChunkedArray<T>,
730    anti: bool,
731    nulls_equal: bool,
732) -> Vec<IdxSize>
733where
734    T: PolarsNumericType,
735    T::Native: TotalHash + TotalEq + DirtyHash + ToTotalOrd,
736    <T::Native as ToTotalOrd>::TotalOrdItem: Send + Sync + Copy + Hash + Eq + DirtyHash + IsNull,
737    <Option<T::Native> as ToTotalOrd>::TotalOrdItem: Send + Sync + DirtyHash + IsNull,
738{
739    let n_threads = RAYON.current_num_threads();
740    let splitted_a = split(left, n_threads);
741    let splitted_b = split(right, n_threads);
742    match (
743        left.null_count(),
744        right.null_count(),
745        left.chunks().len(),
746        right.chunks().len(),
747    ) {
748        (0, 0, 1, 1) => {
749            let keys_a = chunks_as_slices(&splitted_a);
750            let keys_b = chunks_as_slices(&splitted_b);
751            if anti {
752                hash_join_tuples_left_anti(keys_a, keys_b, nulls_equal)
753            } else {
754                hash_join_tuples_left_semi(keys_a, keys_b, nulls_equal)
755            }
756        },
757        (0, 0, _, _) => {
758            let keys_a = chunks_as_slices(&splitted_a);
759            let keys_b = chunks_as_slices(&splitted_b);
760            if anti {
761                hash_join_tuples_left_anti(keys_a, keys_b, nulls_equal)
762            } else {
763                hash_join_tuples_left_semi(keys_a, keys_b, nulls_equal)
764            }
765        },
766        _ => {
767            let keys_a = get_arrays(&splitted_a);
768            let keys_b = get_arrays(&splitted_b);
769            if anti {
770                hash_join_tuples_left_anti(keys_a, keys_b, nulls_equal)
771            } else {
772                hash_join_tuples_left_semi(keys_a, keys_b, nulls_equal)
773            }
774        },
775    }
776}
polars_ops/frame/join/hash_join/single_keys_dispatch.rs

polars_ops/frame/join/hash_join/
single_keys_dispatch.rs
