1use std::fmt::{self, Debug};
29use std::mem::{size_of, size_of_val};
30use std::sync::Arc;
31use std::task::{Context, Poll};
32use std::vec;
33
34use crate::common::SharedMemoryReservation;
35use crate::execution_plan::{boundedness_from_children, emission_type_from_children};
36use crate::joins::stream_join_utils::{
37 PruningJoinHashMap, SortedFilterExpr, StreamJoinMetrics,
38 calculate_filter_expr_intervals, combine_two_batches,
39 convert_sort_expr_with_filter_schema, get_pruning_anti_indices,
40 get_pruning_semi_indices, prepare_sorted_exprs, record_visited_indices,
41};
42use crate::joins::utils::{
43 BatchSplitter, BatchTransformer, ColumnIndex, JoinFilter, JoinHashMapType, JoinOn,
44 JoinOnRef, NoopBatchTransformer, StatefulStreamResult, apply_join_filter_to_indices,
45 build_batch_from_indices, build_join_schema, check_join_is_valid, equal_rows_arr,
46 matchable_join_keys, symmetric_join_output_partitioning, update_hash,
47};
48use crate::projection::{
49 JoinData, ProjectionExec, try_pushdown_through_join_with_column_indices,
50};
51use crate::stream::EmptyRecordBatchStream;
52use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count};
53use crate::{
54 DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties,
55 InputDistributionRequirements, PlanProperties, RecordBatchStream,
56 SendableRecordBatchStream,
57 joins::StreamJoinPartitionMode,
58 metrics::{ExecutionPlanMetricsSet, MetricsSet},
59};
60
61use arrow::array::{
62 ArrowPrimitiveType, NativeAdapter, PrimitiveArray, PrimitiveBuilder, UInt32Array,
63 UInt64Array,
64};
65use arrow::compute::concat_batches;
66use arrow::datatypes::{ArrowNativeType, Schema, SchemaRef};
67use arrow::record_batch::RecordBatch;
68use datafusion_common::hash_utils::create_hashes;
69use datafusion_common::tree_node::TreeNodeRecursion;
70use datafusion_common::utils::bisect;
71use datafusion_common::{
72 HashSet, JoinSide, JoinType, NullEquality, Result, assert_eq_or_internal_err,
73 plan_err,
74};
75use datafusion_execution::TaskContext;
76use datafusion_execution::memory_pool::MemoryConsumer;
77use datafusion_expr::interval_arithmetic::Interval;
78use datafusion_physical_expr::equivalence::join_equivalence_properties;
79use datafusion_physical_expr::intervals::cp_solver::ExprIntervalGraph;
80use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql};
81use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements};
82
83use datafusion_common::hash_utils::RandomState;
84use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays;
85use futures::{Stream, StreamExt, ready};
86
87const HASHMAP_SHRINK_SCALE_FACTOR: usize = 4;
88
89#[derive(Debug, Clone)]
175pub struct SymmetricHashJoinExec {
176 pub(crate) left: Arc<dyn ExecutionPlan>,
178 pub(crate) right: Arc<dyn ExecutionPlan>,
180 pub(crate) on: Vec<(PhysicalExprRef, PhysicalExprRef)>,
182 pub(crate) filter: Option<JoinFilter>,
184 pub(crate) join_type: JoinType,
186 random_state: RandomState,
188 metrics: ExecutionPlanMetricsSet,
190 column_indices: Vec<ColumnIndex>,
192 pub(crate) null_equality: NullEquality,
194 pub(crate) left_sort_exprs: Option<LexOrdering>,
196 pub(crate) right_sort_exprs: Option<LexOrdering>,
198 mode: StreamJoinPartitionMode,
200 cache: Arc<PlanProperties>,
202}
203
204impl SymmetricHashJoinExec {
205 #[expect(clippy::too_many_arguments)]
212 pub fn try_new(
213 left: Arc<dyn ExecutionPlan>,
214 right: Arc<dyn ExecutionPlan>,
215 on: JoinOn,
216 filter: Option<JoinFilter>,
217 join_type: &JoinType,
218 null_equality: NullEquality,
219 left_sort_exprs: Option<LexOrdering>,
220 right_sort_exprs: Option<LexOrdering>,
221 mode: StreamJoinPartitionMode,
222 ) -> Result<Self> {
223 let left_schema = left.schema();
224 let right_schema = right.schema();
225
226 if on.is_empty() {
228 return plan_err!(
229 "On constraints in SymmetricHashJoinExec should be non-empty"
230 );
231 }
232
233 check_join_is_valid(&left_schema, &right_schema, &on)?;
235
236 let (schema, column_indices) =
238 build_join_schema(&left_schema, &right_schema, join_type);
239
240 let random_state = RandomState::with_seed(0);
242 let schema = Arc::new(schema);
243 let cache = Self::compute_properties(&left, &right, schema, *join_type, &on)?;
244 Ok(SymmetricHashJoinExec {
245 left,
246 right,
247 on,
248 filter,
249 join_type: *join_type,
250 random_state,
251 metrics: ExecutionPlanMetricsSet::new(),
252 column_indices,
253 null_equality,
254 left_sort_exprs,
255 right_sort_exprs,
256 mode,
257 cache: Arc::new(cache),
258 })
259 }
260
261 fn compute_properties(
263 left: &Arc<dyn ExecutionPlan>,
264 right: &Arc<dyn ExecutionPlan>,
265 schema: SchemaRef,
266 join_type: JoinType,
267 join_on: JoinOnRef,
268 ) -> Result<PlanProperties> {
269 let eq_properties = join_equivalence_properties(
271 left.equivalence_properties().clone(),
272 right.equivalence_properties().clone(),
273 &join_type,
274 schema,
275 &[false, false],
276 None,
278 join_on,
279 )?;
280
281 let output_partitioning =
282 symmetric_join_output_partitioning(left, right, &join_type)?;
283
284 Ok(PlanProperties::new(
285 eq_properties,
286 output_partitioning,
287 emission_type_from_children([left, right]),
288 boundedness_from_children([left, right]),
289 ))
290 }
291
292 pub fn left(&self) -> &Arc<dyn ExecutionPlan> {
294 &self.left
295 }
296
297 pub fn right(&self) -> &Arc<dyn ExecutionPlan> {
299 &self.right
300 }
301
302 pub fn on(&self) -> &[(PhysicalExprRef, PhysicalExprRef)] {
304 &self.on
305 }
306
307 pub fn filter(&self) -> Option<&JoinFilter> {
309 self.filter.as_ref()
310 }
311
312 pub fn join_type(&self) -> &JoinType {
314 &self.join_type
315 }
316
317 pub fn null_equality(&self) -> NullEquality {
319 self.null_equality
320 }
321
322 pub fn partition_mode(&self) -> StreamJoinPartitionMode {
324 self.mode
325 }
326
327 pub fn left_sort_exprs(&self) -> Option<&LexOrdering> {
329 self.left_sort_exprs.as_ref()
330 }
331
332 pub fn right_sort_exprs(&self) -> Option<&LexOrdering> {
334 self.right_sort_exprs.as_ref()
335 }
336
337 pub fn check_if_order_information_available(&self) -> Result<bool> {
339 if let Some(filter) = self.filter() {
340 let left = self.left();
341 if let Some(left_ordering) = left.output_ordering() {
342 let right = self.right();
343 if let Some(right_ordering) = right.output_ordering() {
344 let left_convertible = convert_sort_expr_with_filter_schema(
345 &JoinSide::Left,
346 filter,
347 &left.schema(),
348 &left_ordering[0],
349 )?
350 .is_some();
351 let right_convertible = convert_sort_expr_with_filter_schema(
352 &JoinSide::Right,
353 filter,
354 &right.schema(),
355 &right_ordering[0],
356 )?
357 .is_some();
358 return Ok(left_convertible && right_convertible);
359 }
360 }
361 }
362 Ok(false)
363 }
364}
365
366impl DisplayAs for SymmetricHashJoinExec {
367 fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
368 match t {
369 DisplayFormatType::Default | DisplayFormatType::Verbose => {
370 let display_filter = self.filter.as_ref().map_or_else(
371 || "".to_string(),
372 |f| format!(", filter={}", f.expression()),
373 );
374 let on = self
375 .on
376 .iter()
377 .map(|(c1, c2)| format!("({c1}, {c2})"))
378 .collect::<Vec<String>>()
379 .join(", ");
380 write!(
381 f,
382 "SymmetricHashJoinExec: mode={:?}, join_type={:?}, on=[{}]{}",
383 self.mode, self.join_type, on, display_filter
384 )
385 }
386 DisplayFormatType::TreeRender => {
387 let on = self
388 .on
389 .iter()
390 .map(|(c1, c2)| {
391 format!("({} = {})", fmt_sql(c1.as_ref()), fmt_sql(c2.as_ref()))
392 })
393 .collect::<Vec<String>>()
394 .join(", ");
395
396 writeln!(f, "mode={:?}", self.mode)?;
397 if *self.join_type() != JoinType::Inner {
398 writeln!(f, "join_type={:?}", self.join_type)?;
399 }
400 writeln!(f, "on={on}")
401 }
402 }
403 }
404}
405
406impl ExecutionPlan for SymmetricHashJoinExec {
407 fn name(&self) -> &'static str {
408 "SymmetricHashJoinExec"
409 }
410
411 fn properties(&self) -> &Arc<PlanProperties> {
412 &self.cache
413 }
414
415 fn required_input_distribution(&self) -> Vec<Distribution> {
416 self.input_distribution_requirements().into_per_child()
417 }
418
419 fn input_distribution_requirements(&self) -> InputDistributionRequirements {
420 match self.mode {
421 StreamJoinPartitionMode::Partitioned => {
422 let (left_expr, right_expr) = self
423 .on
424 .iter()
425 .map(|(l, r)| (Arc::clone(l) as _, Arc::clone(r) as _))
426 .unzip();
427 InputDistributionRequirements::co_partitioned(vec![
428 Distribution::KeyPartitioned(left_expr),
429 Distribution::KeyPartitioned(right_expr),
430 ])
431 }
432 StreamJoinPartitionMode::SinglePartition => {
433 InputDistributionRequirements::new(vec![
434 Distribution::SinglePartition,
435 Distribution::SinglePartition,
436 ])
437 }
438 }
439 }
440
441 fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
442 vec![
443 self.left_sort_exprs
444 .as_ref()
445 .map(|e| OrderingRequirements::from(e.clone())),
446 self.right_sort_exprs
447 .as_ref()
448 .map(|e| OrderingRequirements::from(e.clone())),
449 ]
450 }
451
452 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
453 vec![&self.left, &self.right]
454 }
455
456 fn apply_expressions(
457 &self,
458 f: &mut dyn FnMut(&Arc<dyn crate::PhysicalExpr>) -> Result<TreeNodeRecursion>,
459 ) -> Result<TreeNodeRecursion> {
460 let join_keys = self.on.iter().flat_map(|(left, right)| [left, right]);
461 let filter = self.filter.iter().map(|filter| filter.expression());
462 crate::apply_expression_roots(join_keys.chain(filter), f)
463 }
464
465 fn replace_children(
466 self: Arc<Self>,
467 mut children: Vec<Arc<dyn ExecutionPlan>>,
468 options: ReplaceChildrenOptions,
469 ) -> Result<Arc<dyn ExecutionPlan>> {
470 validate_child_count!(self, children);
471 match options.children_properties {
472 ChildrenPropertiesMode::Keep => {
473 let left = children.swap_remove(0);
474 let right = children.swap_remove(0);
475 Ok(Arc::new(Self {
476 left,
477 right,
478 metrics: ExecutionPlanMetricsSet::new(),
479 ..Self::clone(&*self)
480 }))
481 }
482 ChildrenPropertiesMode::Recompute => {
483 Ok(Arc::new(SymmetricHashJoinExec::try_new(
484 Arc::clone(&children[0]),
485 Arc::clone(&children[1]),
486 self.on.clone(),
487 self.filter.clone(),
488 &self.join_type,
489 self.null_equality,
490 self.left_sort_exprs.clone(),
491 self.right_sort_exprs.clone(),
492 self.mode,
493 )?))
494 }
495 }
496 }
497
498 fn with_new_children(
499 self: Arc<Self>,
500 children: Vec<Arc<dyn ExecutionPlan>>,
501 ) -> Result<Arc<dyn ExecutionPlan>> {
502 self.replace_children(
503 children,
504 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
505 )
506 }
507
508 fn with_new_children_and_same_properties(
509 self: Arc<Self>,
510 children: Vec<Arc<dyn ExecutionPlan>>,
511 ) -> Result<Arc<dyn ExecutionPlan>> {
512 self.replace_children(
513 children,
514 ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
515 )
516 }
517
518 fn metrics(&self) -> Option<MetricsSet> {
519 Some(self.metrics.clone_inner())
520 }
521
522 fn execute(
523 &self,
524 partition: usize,
525 context: Arc<TaskContext>,
526 ) -> Result<SendableRecordBatchStream> {
527 let left_partitions = self.left.output_partitioning().partition_count();
528 let right_partitions = self.right.output_partitioning().partition_count();
529 assert_eq_or_internal_err!(
530 left_partitions,
531 right_partitions,
532 "Invalid SymmetricHashJoinExec, partition count mismatch {left_partitions}!={right_partitions},\
533 consider using RepartitionExec"
534 );
535 let (left_sorted_filter_expr, right_sorted_filter_expr, graph) = match (
538 self.left_sort_exprs(),
539 self.right_sort_exprs(),
540 &self.filter,
541 ) {
542 (Some(left_sort_exprs), Some(right_sort_exprs), Some(filter)) => {
543 let (left, right, graph) = prepare_sorted_exprs(
544 filter,
545 &self.left,
546 &self.right,
547 left_sort_exprs,
548 right_sort_exprs,
549 )?;
550 (Some(left), Some(right), Some(graph))
551 }
552 _ => (None, None, None),
555 };
556
557 let (on_left, on_right) = self.on.iter().cloned().unzip();
558
559 let left_side_joiner =
560 OneSideHashJoiner::new(JoinSide::Left, on_left, self.left.schema());
561 let right_side_joiner =
562 OneSideHashJoiner::new(JoinSide::Right, on_right, self.right.schema());
563
564 let left_stream = self.left.execute(partition, Arc::clone(&context))?;
565
566 let right_stream = self.right.execute(partition, Arc::clone(&context))?;
567
568 let batch_size = context.session_config().batch_size();
569 let enforce_batch_size_in_joins =
570 context.session_config().enforce_batch_size_in_joins();
571
572 let reservation = Arc::new(
573 MemoryConsumer::new(format!("SymmetricHashJoinStream[{partition}]"))
574 .register(context.memory_pool()),
575 );
576 if let Some(g) = graph.as_ref() {
577 reservation.try_grow(g.size())?;
578 }
579
580 if enforce_batch_size_in_joins {
581 Ok(Box::pin(SymmetricHashJoinStream {
582 left_stream,
583 right_stream,
584 schema: self.schema(),
585 filter: self.filter.clone(),
586 join_type: self.join_type,
587 random_state: self.random_state.clone(),
588 left: left_side_joiner,
589 right: right_side_joiner,
590 column_indices: self.column_indices.clone(),
591 metrics: StreamJoinMetrics::new(partition, &self.metrics),
592 graph,
593 left_sorted_filter_expr,
594 right_sorted_filter_expr,
595 null_equality: self.null_equality,
596 state: SHJStreamState::PullRight,
597 reservation,
598 batch_transformer: BatchSplitter::new(batch_size),
599 }))
600 } else {
601 Ok(Box::pin(SymmetricHashJoinStream {
602 left_stream,
603 right_stream,
604 schema: self.schema(),
605 filter: self.filter.clone(),
606 join_type: self.join_type,
607 random_state: self.random_state.clone(),
608 left: left_side_joiner,
609 right: right_side_joiner,
610 column_indices: self.column_indices.clone(),
611 metrics: StreamJoinMetrics::new(partition, &self.metrics),
612 graph,
613 left_sorted_filter_expr,
614 right_sorted_filter_expr,
615 null_equality: self.null_equality,
616 state: SHJStreamState::PullRight,
617 reservation,
618 batch_transformer: NoopBatchTransformer::new(),
619 }))
620 }
621 }
622
623 fn try_swapping_with_projection(
627 &self,
628 projection: &ProjectionExec,
629 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
630 let schema = self.schema();
631 if let Some(JoinData {
632 projected_left_child,
633 projected_right_child,
634 join_filter,
635 join_on,
636 }) = try_pushdown_through_join_with_column_indices(
637 projection,
638 self.left(),
639 self.right(),
640 self.on(),
641 &schema,
642 self.filter(),
643 self.column_indices.as_slice(),
644 )? {
645 SymmetricHashJoinExec::try_new(
646 Arc::new(projected_left_child),
647 Arc::new(projected_right_child),
648 join_on,
649 join_filter,
650 self.join_type(),
651 self.null_equality(),
652 self.right().output_ordering().cloned(),
653 self.left().output_ordering().cloned(),
654 self.partition_mode(),
655 )
656 .map(|e| Some(Arc::new(e) as _))
657 } else {
658 Ok(None)
659 }
660 }
661
662 #[cfg(feature = "proto")]
663 fn try_to_proto(
664 &self,
665 ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
666 ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
667 use datafusion_proto_models::protobuf;
668
669 let left = ctx.encode_child(self.left())?;
670 let right = ctx.encode_child(self.right())?;
671 let on = self
672 .on()
673 .iter()
674 .map(|(left, right)| {
675 Ok(protobuf::JoinOn {
676 left: Some(ctx.encode_expr(left)?),
677 right: Some(ctx.encode_expr(right)?),
678 })
679 })
680 .collect::<Result<Vec<_>>>()?;
681
682 let join_type = match self.join_type() {
683 JoinType::Inner => protobuf::JoinType::Inner,
684 JoinType::Left => protobuf::JoinType::Left,
685 JoinType::Right => protobuf::JoinType::Right,
686 JoinType::Full => protobuf::JoinType::Full,
687 JoinType::LeftSemi => protobuf::JoinType::Leftsemi,
688 JoinType::RightSemi => protobuf::JoinType::Rightsemi,
689 JoinType::LeftAnti => protobuf::JoinType::Leftanti,
690 JoinType::RightAnti => protobuf::JoinType::Rightanti,
691 JoinType::LeftMark => protobuf::JoinType::Leftmark,
692 JoinType::RightMark => protobuf::JoinType::Rightmark,
693 };
694 let null_equality = match self.null_equality() {
695 NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing,
696 NullEquality::NullEqualsNull => protobuf::NullEquality::NullEqualsNull,
697 };
698 let partition_mode = match self.partition_mode() {
699 StreamJoinPartitionMode::SinglePartition => {
700 protobuf::StreamPartitionMode::SinglePartition
701 }
702 StreamJoinPartitionMode::Partitioned => {
703 protobuf::StreamPartitionMode::PartitionedExec
704 }
705 };
706 let filter = self
707 .filter()
708 .map(|filter| -> Result<protobuf::JoinFilter> {
709 let expression = ctx.encode_expr(filter.expression())?;
710 let column_indices = filter
711 .column_indices()
712 .iter()
713 .map(|column_index| {
714 let side = match column_index.side {
715 JoinSide::Left => protobuf::JoinSide::LeftSide,
716 JoinSide::Right => protobuf::JoinSide::RightSide,
717 JoinSide::None => protobuf::JoinSide::None,
718 };
719 protobuf::ColumnIndex {
720 index: column_index.index as u32,
721 side: side.into(),
722 }
723 })
724 .collect();
725 Ok(protobuf::JoinFilter {
726 expression: Some(expression),
727 column_indices,
728 schema: Some(filter.schema().as_ref().try_into()?),
729 })
730 })
731 .transpose()?;
732 let expr_ctx = ctx.expr_ctx();
733 let left_sort_exprs =
734 datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto(
735 self.left_sort_exprs(),
736 &expr_ctx,
737 )?;
738 let right_sort_exprs =
739 datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto(
740 self.right_sort_exprs(),
741 &expr_ctx,
742 )?;
743
744 Ok(Some(protobuf::PhysicalPlanNode {
745 physical_plan_type: Some(
746 protobuf::physical_plan_node::PhysicalPlanType::SymmetricHashJoin(
747 Box::new(protobuf::SymmetricHashJoinExecNode {
748 left: Some(Box::new(left)),
749 right: Some(Box::new(right)),
750 on,
751 join_type: join_type.into(),
752 partition_mode: partition_mode.into(),
753 null_equality: null_equality.into(),
754 filter,
755 left_sort_exprs,
756 right_sort_exprs,
757 }),
758 ),
759 ),
760 }))
761 }
762}
763
764#[cfg(feature = "proto")]
765impl SymmetricHashJoinExec {
766 pub fn try_from_proto(
772 node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
773 ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
774 ) -> Result<Arc<dyn ExecutionPlan>> {
775 use datafusion_common::internal_datafusion_err;
776 use datafusion_proto_models::protobuf;
777
778 let sym_join = crate::expect_plan_variant!(
779 node,
780 protobuf::physical_plan_node::PhysicalPlanType::SymmetricHashJoin,
781 "SymmetricHashJoinExec",
782 );
783 let left = ctx.decode_required_child(
784 sym_join.left.as_deref(),
785 "SymmetricHashJoinExec",
786 "left",
787 )?;
788 let right = ctx.decode_required_child(
789 sym_join.right.as_deref(),
790 "SymmetricHashJoinExec",
791 "right",
792 )?;
793 let left_schema = left.schema();
794 let right_schema = right.schema();
795 let on = sym_join
796 .on
797 .iter()
798 .map(|columns| {
799 let left = ctx.decode_required_expr(
800 columns.left.as_ref(),
801 left_schema.as_ref(),
802 "SymmetricHashJoinExec",
803 "on.left",
804 )?;
805 let right = ctx.decode_required_expr(
806 columns.right.as_ref(),
807 right_schema.as_ref(),
808 "SymmetricHashJoinExec",
809 "on.right",
810 )?;
811 Ok((left, right))
812 })
813 .collect::<Result<JoinOn>>()?;
814
815 let join_type =
816 match protobuf::JoinType::try_from(sym_join.join_type).map_err(|_| {
817 internal_datafusion_err!(
818 "SymmetricHashJoinExec: unknown JoinType {}",
819 sym_join.join_type
820 )
821 })? {
822 protobuf::JoinType::Inner => JoinType::Inner,
823 protobuf::JoinType::Left => JoinType::Left,
824 protobuf::JoinType::Right => JoinType::Right,
825 protobuf::JoinType::Full => JoinType::Full,
826 protobuf::JoinType::Leftsemi => JoinType::LeftSemi,
827 protobuf::JoinType::Rightsemi => JoinType::RightSemi,
828 protobuf::JoinType::Leftanti => JoinType::LeftAnti,
829 protobuf::JoinType::Rightanti => JoinType::RightAnti,
830 protobuf::JoinType::Leftmark => JoinType::LeftMark,
831 protobuf::JoinType::Rightmark => JoinType::RightMark,
832 };
833 let null_equality = match protobuf::NullEquality::try_from(sym_join.null_equality)
834 .map_err(|_| {
835 internal_datafusion_err!(
836 "SymmetricHashJoinExec: unknown NullEquality {}",
837 sym_join.null_equality
838 )
839 })? {
840 protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing,
841 protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull,
842 };
843 let partition_mode =
844 match protobuf::StreamPartitionMode::try_from(sym_join.partition_mode)
845 .map_err(|_| {
846 internal_datafusion_err!(
847 "SymmetricHashJoinExec: unknown StreamPartitionMode {}",
848 sym_join.partition_mode
849 )
850 })? {
851 protobuf::StreamPartitionMode::SinglePartition => {
852 StreamJoinPartitionMode::SinglePartition
853 }
854 protobuf::StreamPartitionMode::PartitionedExec => {
855 StreamJoinPartitionMode::Partitioned
856 }
857 };
858 let filter = sym_join
859 .filter
860 .as_ref()
861 .map(|filter| -> Result<JoinFilter> {
862 let schema: Schema = filter
863 .schema
864 .as_ref()
865 .ok_or_else(|| {
866 internal_datafusion_err!(
867 "SymmetricHashJoinExec: JoinFilter missing schema"
868 )
869 })?
870 .try_into()?;
871 let expression = ctx.decode_required_expr(
872 filter.expression.as_ref(),
873 &schema,
874 "SymmetricHashJoinExec",
875 "filter.expression",
876 )?;
877 let column_indices = filter
878 .column_indices
879 .iter()
880 .map(|column_index| {
881 let side = protobuf::JoinSide::try_from(column_index.side)
882 .map_err(|_| {
883 internal_datafusion_err!(
884 "SymmetricHashJoinExec: unknown JoinSide {}",
885 column_index.side
886 )
887 })?;
888 let side = match side {
889 protobuf::JoinSide::LeftSide => JoinSide::Left,
890 protobuf::JoinSide::RightSide => JoinSide::Right,
891 protobuf::JoinSide::None => JoinSide::None,
892 };
893 Ok(ColumnIndex {
894 index: column_index.index as usize,
895 side,
896 })
897 })
898 .collect::<Result<Vec<_>>>()?;
899 Ok(JoinFilter::new(
900 expression,
901 column_indices,
902 Arc::new(schema),
903 ))
904 })
905 .transpose()?;
906 let left_sort_exprs =
907 datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto(
908 &sym_join.left_sort_exprs,
909 &ctx.expr_ctx(left_schema.as_ref()),
910 )?;
911 let right_sort_exprs =
912 datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto(
913 &sym_join.right_sort_exprs,
914 &ctx.expr_ctx(right_schema.as_ref()),
915 )?;
916
917 Self::try_new(
918 left,
919 right,
920 on,
921 filter,
922 &join_type,
923 null_equality,
924 left_sort_exprs,
925 right_sort_exprs,
926 partition_mode,
927 )
928 .map(|exec| Arc::new(exec) as _)
929 }
930}
931
932struct SymmetricHashJoinStream<T> {
934 left_stream: SendableRecordBatchStream,
936 right_stream: SendableRecordBatchStream,
937 schema: Arc<Schema>,
939 filter: Option<JoinFilter>,
941 join_type: JoinType,
943 left: OneSideHashJoiner,
945 right: OneSideHashJoiner,
947 column_indices: Vec<ColumnIndex>,
949 graph: Option<ExprIntervalGraph>,
951 left_sorted_filter_expr: Option<SortedFilterExpr>,
953 right_sorted_filter_expr: Option<SortedFilterExpr>,
955 random_state: RandomState,
957 null_equality: NullEquality,
959 metrics: StreamJoinMetrics,
961 reservation: SharedMemoryReservation,
963 state: SHJStreamState,
965 batch_transformer: T,
967}
968
969impl<T: BatchTransformer + Unpin + Send> RecordBatchStream
970 for SymmetricHashJoinStream<T>
971{
972 fn schema(&self) -> SchemaRef {
973 Arc::clone(&self.schema)
974 }
975}
976
977impl<T: BatchTransformer + Unpin + Send> Stream for SymmetricHashJoinStream<T> {
978 type Item = Result<RecordBatch>;
979
980 fn poll_next(
981 mut self: std::pin::Pin<&mut Self>,
982 cx: &mut Context<'_>,
983 ) -> Poll<Option<Self::Item>> {
984 self.poll_next_impl(cx)
985 }
986}
987
988fn determine_prune_length(
1007 buffer: &RecordBatch,
1008 build_side_filter_expr: &SortedFilterExpr,
1009) -> Result<usize> {
1010 let origin_sorted_expr = build_side_filter_expr.origin_sorted_expr();
1011 let interval = build_side_filter_expr.interval();
1012 let batch_arr = origin_sorted_expr
1014 .expr
1015 .evaluate(buffer)?
1016 .into_array(buffer.num_rows())?;
1017
1018 let target = if origin_sorted_expr.options.descending {
1020 interval.upper().clone()
1021 } else {
1022 interval.lower().clone()
1023 };
1024
1025 bisect::<true>(&[batch_arr], &[target], &[origin_sorted_expr.options])
1027}
1028
1029fn need_to_produce_result_in_final(build_side: JoinSide, join_type: JoinType) -> bool {
1044 if build_side == JoinSide::Left {
1045 matches!(
1046 join_type,
1047 JoinType::Left
1048 | JoinType::LeftAnti
1049 | JoinType::Full
1050 | JoinType::LeftSemi
1051 | JoinType::LeftMark
1052 )
1053 } else {
1054 matches!(
1055 join_type,
1056 JoinType::Right
1057 | JoinType::RightAnti
1058 | JoinType::Full
1059 | JoinType::RightSemi
1060 | JoinType::RightMark
1061 )
1062 }
1063}
1064
1065fn calculate_indices_by_join_type<L: ArrowPrimitiveType, R: ArrowPrimitiveType>(
1082 build_side: JoinSide,
1083 prune_length: usize,
1084 visited_rows: &HashSet<usize>,
1085 deleted_offset: usize,
1086 join_type: JoinType,
1087) -> Result<(PrimitiveArray<L>, PrimitiveArray<R>)>
1088where
1089 NativeAdapter<L>: From<<L as ArrowPrimitiveType>::Native>,
1090{
1091 let result = match (build_side, join_type) {
1093 (JoinSide::Left, JoinType::LeftMark) => {
1109 let build_indices = (0..prune_length)
1110 .map(L::Native::from_usize)
1111 .collect::<PrimitiveArray<L>>();
1112 let probe_indices = (0..prune_length)
1113 .map(|idx| {
1114 visited_rows
1116 .contains(&(idx + deleted_offset))
1117 .then_some(R::Native::from_usize(0).unwrap())
1118 })
1119 .collect();
1120 (build_indices, probe_indices)
1121 }
1122 (JoinSide::Right, JoinType::RightMark) => {
1123 let build_indices = (0..prune_length)
1124 .map(L::Native::from_usize)
1125 .collect::<PrimitiveArray<L>>();
1126 let probe_indices = (0..prune_length)
1127 .map(|idx| {
1128 visited_rows
1130 .contains(&(idx + deleted_offset))
1131 .then_some(R::Native::from_usize(0).unwrap())
1132 })
1133 .collect();
1134 (build_indices, probe_indices)
1135 }
1136 (JoinSide::Left, JoinType::Left | JoinType::LeftAnti)
1138 | (JoinSide::Right, JoinType::Right | JoinType::RightAnti)
1139 | (_, JoinType::Full) => {
1140 let build_unmatched_indices =
1141 get_pruning_anti_indices(prune_length, deleted_offset, visited_rows);
1142 let mut builder =
1143 PrimitiveBuilder::<R>::with_capacity(build_unmatched_indices.len());
1144 builder.append_nulls(build_unmatched_indices.len());
1145 let probe_indices = builder.finish();
1146 (build_unmatched_indices, probe_indices)
1147 }
1148 (JoinSide::Left, JoinType::LeftSemi) | (JoinSide::Right, JoinType::RightSemi) => {
1150 let build_unmatched_indices =
1151 get_pruning_semi_indices(prune_length, deleted_offset, visited_rows);
1152 let mut builder =
1153 PrimitiveBuilder::<R>::with_capacity(build_unmatched_indices.len());
1154 builder.append_nulls(build_unmatched_indices.len());
1155 let probe_indices = builder.finish();
1156 (build_unmatched_indices, probe_indices)
1157 }
1158 _ => unreachable!(),
1160 };
1161 Ok(result)
1162}
1163
1164pub(crate) fn build_side_determined_results(
1182 build_hash_joiner: &OneSideHashJoiner,
1183 output_schema: &SchemaRef,
1184 prune_length: usize,
1185 probe_schema: SchemaRef,
1186 join_type: JoinType,
1187 column_indices: &[ColumnIndex],
1188) -> Result<Option<RecordBatch>> {
1189 if prune_length > 0
1191 && need_to_produce_result_in_final(build_hash_joiner.build_side, join_type)
1192 {
1193 let (build_indices, probe_indices) = calculate_indices_by_join_type(
1195 build_hash_joiner.build_side,
1196 prune_length,
1197 &build_hash_joiner.visited_rows,
1198 build_hash_joiner.deleted_offset,
1199 join_type,
1200 )?;
1201
1202 let empty_probe_batch = RecordBatch::new_empty(probe_schema);
1204 build_batch_from_indices(
1206 output_schema.as_ref(),
1207 &build_hash_joiner.input_buffer,
1208 &empty_probe_batch,
1209 &build_indices,
1210 &probe_indices,
1211 column_indices,
1212 build_hash_joiner.build_side,
1213 join_type,
1214 )
1215 .map(|batch| (batch.num_rows() > 0).then_some(batch))
1216 } else {
1217 Ok(None)
1219 }
1220}
1221
1222#[expect(clippy::too_many_arguments)]
1242pub(crate) fn join_with_probe_batch(
1243 build_hash_joiner: &mut OneSideHashJoiner,
1244 probe_hash_joiner: &mut OneSideHashJoiner,
1245 schema: &SchemaRef,
1246 join_type: JoinType,
1247 filter: Option<&JoinFilter>,
1248 probe_batch: &RecordBatch,
1249 column_indices: &[ColumnIndex],
1250 random_state: &RandomState,
1251 null_equality: NullEquality,
1252) -> Result<Option<RecordBatch>> {
1253 if build_hash_joiner.input_buffer.num_rows() == 0 || probe_batch.num_rows() == 0 {
1254 return Ok(None);
1255 }
1256 let (build_indices, probe_indices) = lookup_join_hashmap(
1257 &build_hash_joiner.hashmap,
1258 &build_hash_joiner.input_buffer,
1259 probe_batch,
1260 &build_hash_joiner.on,
1261 &probe_hash_joiner.on,
1262 random_state,
1263 null_equality,
1264 &mut build_hash_joiner.hashes_buffer,
1265 Some(build_hash_joiner.deleted_offset),
1266 )?;
1267
1268 let (build_indices, probe_indices) = if let Some(filter) = filter {
1269 apply_join_filter_to_indices(
1270 &build_hash_joiner.input_buffer,
1271 probe_batch,
1272 build_indices,
1273 probe_indices,
1274 filter,
1275 build_hash_joiner.build_side,
1276 None,
1277 join_type,
1278 )?
1279 } else {
1280 (build_indices, probe_indices)
1281 };
1282
1283 if need_to_produce_result_in_final(build_hash_joiner.build_side, join_type) {
1284 record_visited_indices(
1285 &mut build_hash_joiner.visited_rows,
1286 build_hash_joiner.deleted_offset,
1287 &build_indices,
1288 );
1289 }
1290 if need_to_produce_result_in_final(build_hash_joiner.build_side.negate(), join_type) {
1291 record_visited_indices(
1292 &mut probe_hash_joiner.visited_rows,
1293 probe_hash_joiner.offset,
1294 &probe_indices,
1295 );
1296 }
1297 if matches!(
1298 join_type,
1299 JoinType::LeftAnti
1300 | JoinType::RightAnti
1301 | JoinType::LeftSemi
1302 | JoinType::LeftMark
1303 | JoinType::RightSemi
1304 | JoinType::RightMark
1305 ) {
1306 Ok(None)
1307 } else {
1308 build_batch_from_indices(
1309 schema,
1310 &build_hash_joiner.input_buffer,
1311 probe_batch,
1312 &build_indices,
1313 &probe_indices,
1314 column_indices,
1315 build_hash_joiner.build_side,
1316 join_type,
1317 )
1318 .map(|batch| (batch.num_rows() > 0).then_some(batch))
1319 }
1320}
1321
1322#[expect(clippy::too_many_arguments)]
1342fn lookup_join_hashmap(
1343 build_hashmap: &PruningJoinHashMap,
1344 build_batch: &RecordBatch,
1345 probe_batch: &RecordBatch,
1346 build_on: &[PhysicalExprRef],
1347 probe_on: &[PhysicalExprRef],
1348 random_state: &RandomState,
1349 null_equality: NullEquality,
1350 hashes_buffer: &mut Vec<u64>,
1351 deleted_offset: Option<usize>,
1352) -> Result<(UInt64Array, UInt32Array)> {
1353 let keys_values = evaluate_expressions_to_arrays(probe_on, probe_batch)?;
1354 let build_join_values = evaluate_expressions_to_arrays(build_on, build_batch)?;
1355
1356 hashes_buffer.clear();
1357 hashes_buffer.resize(probe_batch.num_rows(), 0);
1358 let hash_values = create_hashes(&keys_values, random_state, hashes_buffer)?;
1359
1360 let valid_keys = matchable_join_keys(&keys_values, null_equality);
1394 let (mut matched_probe, mut matched_build) = build_hashmap.get_matched_indices(
1395 Box::new(
1396 hash_values
1397 .iter()
1398 .enumerate()
1399 .filter(|(i, _)| {
1400 valid_keys.as_ref().is_none_or(|valid| valid.is_valid(*i))
1401 })
1402 .rev(),
1403 ),
1404 deleted_offset,
1405 );
1406
1407 matched_probe.reverse();
1408 matched_build.reverse();
1409
1410 let build_indices: UInt64Array = matched_build.into();
1411 let probe_indices: UInt32Array = matched_probe.into();
1412
1413 let (build_indices, probe_indices) = equal_rows_arr(
1414 &build_indices,
1415 &probe_indices,
1416 &build_join_values,
1417 &keys_values,
1418 null_equality,
1419 )?;
1420
1421 Ok((build_indices, probe_indices))
1422}
1423
1424pub struct OneSideHashJoiner {
1425 build_side: JoinSide,
1427 pub input_buffer: RecordBatch,
1429 pub(crate) on: Vec<PhysicalExprRef>,
1431 pub(crate) hashmap: PruningJoinHashMap,
1433 pub(crate) hashes_buffer: Vec<u64>,
1435 pub(crate) visited_rows: HashSet<usize>,
1437 pub(crate) offset: usize,
1439 pub(crate) deleted_offset: usize,
1441}
1442
1443impl OneSideHashJoiner {
1444 pub fn size(&self) -> usize {
1445 let mut size = 0;
1446 size += size_of_val(self);
1447 size += size_of_val(&self.build_side);
1448 size += self.input_buffer.get_array_memory_size();
1449 size += size_of_val(&self.on);
1450 size += self.hashmap.size();
1451 size += self.hashes_buffer.capacity() * size_of::<u64>();
1452 size += self.visited_rows.capacity() * size_of::<usize>();
1453 size += size_of_val(&self.offset);
1454 size += size_of_val(&self.deleted_offset);
1455 size
1456 }
1457 pub fn new(
1458 build_side: JoinSide,
1459 on: Vec<PhysicalExprRef>,
1460 schema: SchemaRef,
1461 ) -> Self {
1462 Self {
1463 build_side,
1464 input_buffer: RecordBatch::new_empty(schema),
1465 on,
1466 hashmap: PruningJoinHashMap::with_capacity(0),
1467 hashes_buffer: vec![],
1468 visited_rows: HashSet::new(),
1469 offset: 0,
1470 deleted_offset: 0,
1471 }
1472 }
1473
1474 pub(crate) fn update_internal_state(
1486 &mut self,
1487 batch: &RecordBatch,
1488 random_state: &RandomState,
1489 null_equality: NullEquality,
1490 ) -> Result<()> {
1491 self.input_buffer = concat_batches(&batch.schema(), [&self.input_buffer, batch])?;
1493 self.hashes_buffer.resize(batch.num_rows(), 0);
1495 update_hash(
1498 &self.on,
1499 batch,
1500 &mut self.hashmap,
1501 self.offset,
1502 random_state,
1503 &mut self.hashes_buffer,
1504 self.deleted_offset,
1505 false,
1506 null_equality,
1507 )?;
1508 Ok(())
1509 }
1510
1511 pub(crate) fn calculate_prune_length_with_probe_batch(
1523 &mut self,
1524 build_side_sorted_filter_expr: &mut SortedFilterExpr,
1525 probe_side_sorted_filter_expr: &mut SortedFilterExpr,
1526 graph: &mut ExprIntervalGraph,
1527 ) -> Result<usize> {
1528 if self.input_buffer.num_rows() == 0 {
1530 return Ok(0);
1531 }
1532 let mut filter_intervals = vec![];
1535 for expr in [
1536 &build_side_sorted_filter_expr,
1537 &probe_side_sorted_filter_expr,
1538 ] {
1539 filter_intervals.push((expr.node_index(), expr.interval().clone()))
1540 }
1541 graph.update_ranges(&mut filter_intervals, Interval::TRUE)?;
1543 let calculated_build_side_interval = filter_intervals.remove(0).1;
1545 if calculated_build_side_interval.eq(build_side_sorted_filter_expr.interval()) {
1547 return Ok(0);
1548 }
1549 build_side_sorted_filter_expr.set_interval(calculated_build_side_interval);
1551
1552 determine_prune_length(&self.input_buffer, build_side_sorted_filter_expr)
1553 }
1554
1555 pub(crate) fn prune_internal_state(&mut self, prune_length: usize) -> Result<()> {
1556 self.hashmap.prune_hash_values(
1558 prune_length,
1559 self.deleted_offset as u64,
1560 HASHMAP_SHRINK_SCALE_FACTOR,
1561 );
1562 for row in self.deleted_offset..(self.deleted_offset + prune_length) {
1564 self.visited_rows.remove(&row);
1565 }
1566 self.input_buffer = self
1568 .input_buffer
1569 .slice(prune_length, self.input_buffer.num_rows() - prune_length);
1570 self.deleted_offset += prune_length;
1572 Ok(())
1573 }
1574}
1575
1576impl<T: BatchTransformer> SymmetricHashJoinStream<T> {
1621 fn poll_next_impl(
1635 &mut self,
1636 cx: &mut Context<'_>,
1637 ) -> Poll<Option<Result<RecordBatch>>> {
1638 loop {
1639 match self.batch_transformer.next() {
1640 None => {
1641 let result = match self.state() {
1642 SHJStreamState::PullRight => {
1643 ready!(self.fetch_next_from_right_stream(cx))
1644 }
1645 SHJStreamState::PullLeft => {
1646 ready!(self.fetch_next_from_left_stream(cx))
1647 }
1648 SHJStreamState::RightExhausted => {
1649 ready!(self.handle_right_stream_end(cx))
1650 }
1651 SHJStreamState::LeftExhausted => {
1652 ready!(self.handle_left_stream_end(cx))
1653 }
1654 SHJStreamState::BothExhausted {
1655 final_result: false,
1656 } => self.prepare_for_final_results_after_exhaustion(),
1657 SHJStreamState::BothExhausted { final_result: true } => {
1658 return Poll::Ready(None);
1659 }
1660 };
1661
1662 match result? {
1663 StatefulStreamResult::Ready(None) => {
1664 return Poll::Ready(None);
1665 }
1666 StatefulStreamResult::Ready(Some(batch)) => {
1667 self.batch_transformer.set_batch(batch);
1668 }
1669 _ => {}
1670 }
1671 }
1672 Some((batch, _)) => {
1673 return self
1674 .metrics
1675 .baseline_metrics
1676 .record_poll(Poll::Ready(Some(Ok(batch))));
1677 }
1678 }
1679 }
1680 }
1681
1682 fn cleanup_depleted_right_stream(&mut self) {
1684 let right_schema = self.right_stream.schema();
1685 self.right_stream = Box::pin(EmptyRecordBatchStream::new(right_schema));
1686 }
1687
1688 fn cleanup_depleted_left_stream(&mut self) {
1690 let left_schema = self.left_stream.schema();
1691 self.left_stream = Box::pin(EmptyRecordBatchStream::new(left_schema));
1692 }
1693
1694 fn fetch_next_from_right_stream(
1704 &mut self,
1705 cx: &mut Context<'_>,
1706 ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
1707 match ready!(self.right_stream().poll_next_unpin(cx)) {
1708 Some(Ok(batch)) => {
1709 if batch.num_rows() == 0 {
1710 return Poll::Ready(Ok(StatefulStreamResult::Continue));
1711 }
1712 self.set_state(SHJStreamState::PullLeft);
1713 Poll::Ready(self.process_batch_from_right(&batch))
1714 }
1715 Some(Err(e)) => Poll::Ready(Err(e)),
1716 None => {
1717 self.cleanup_depleted_right_stream();
1718 self.set_state(SHJStreamState::RightExhausted);
1719 Poll::Ready(Ok(StatefulStreamResult::Continue))
1720 }
1721 }
1722 }
1723
1724 fn fetch_next_from_left_stream(
1734 &mut self,
1735 cx: &mut Context<'_>,
1736 ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
1737 match ready!(self.left_stream().poll_next_unpin(cx)) {
1738 Some(Ok(batch)) => {
1739 if batch.num_rows() == 0 {
1740 return Poll::Ready(Ok(StatefulStreamResult::Continue));
1741 }
1742 self.set_state(SHJStreamState::PullRight);
1743 Poll::Ready(self.process_batch_from_left(&batch))
1744 }
1745 Some(Err(e)) => Poll::Ready(Err(e)),
1746 None => {
1747 self.cleanup_depleted_left_stream();
1748 self.set_state(SHJStreamState::LeftExhausted);
1749 Poll::Ready(Ok(StatefulStreamResult::Continue))
1750 }
1751 }
1752 }
1753
1754 fn handle_right_stream_end(
1765 &mut self,
1766 cx: &mut Context<'_>,
1767 ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
1768 match ready!(self.left_stream().poll_next_unpin(cx)) {
1769 Some(Ok(batch)) => {
1770 if batch.num_rows() == 0 {
1771 return Poll::Ready(Ok(StatefulStreamResult::Continue));
1772 }
1773 Poll::Ready(self.process_batch_after_right_end(&batch))
1774 }
1775 Some(Err(e)) => Poll::Ready(Err(e)),
1776 None => {
1777 self.cleanup_depleted_left_stream();
1778 self.set_state(SHJStreamState::BothExhausted {
1779 final_result: false,
1780 });
1781 Poll::Ready(Ok(StatefulStreamResult::Continue))
1782 }
1783 }
1784 }
1785
1786 fn handle_left_stream_end(
1797 &mut self,
1798 cx: &mut Context<'_>,
1799 ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
1800 match ready!(self.right_stream().poll_next_unpin(cx)) {
1801 Some(Ok(batch)) => {
1802 if batch.num_rows() == 0 {
1803 return Poll::Ready(Ok(StatefulStreamResult::Continue));
1804 }
1805 Poll::Ready(self.process_batch_after_left_end(&batch))
1806 }
1807 Some(Err(e)) => Poll::Ready(Err(e)),
1808 None => {
1809 self.cleanup_depleted_right_stream();
1810 self.set_state(SHJStreamState::BothExhausted {
1811 final_result: false,
1812 });
1813 Poll::Ready(Ok(StatefulStreamResult::Continue))
1814 }
1815 }
1816 }
1817
1818 fn prepare_for_final_results_after_exhaustion(
1828 &mut self,
1829 ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
1830 self.set_state(SHJStreamState::BothExhausted { final_result: true });
1831 self.process_batches_before_finalization()
1832 }
1833
1834 fn process_batch_from_right(
1835 &mut self,
1836 batch: &RecordBatch,
1837 ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
1838 self.perform_join_for_given_side(batch, JoinSide::Right)
1839 .map(|maybe_batch| {
1840 if maybe_batch.is_some() {
1841 StatefulStreamResult::Ready(maybe_batch)
1842 } else {
1843 StatefulStreamResult::Continue
1844 }
1845 })
1846 }
1847
1848 fn process_batch_from_left(
1849 &mut self,
1850 batch: &RecordBatch,
1851 ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
1852 self.perform_join_for_given_side(batch, JoinSide::Left)
1853 .map(|maybe_batch| {
1854 if maybe_batch.is_some() {
1855 StatefulStreamResult::Ready(maybe_batch)
1856 } else {
1857 StatefulStreamResult::Continue
1858 }
1859 })
1860 }
1861
1862 fn process_batch_after_left_end(
1863 &mut self,
1864 right_batch: &RecordBatch,
1865 ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
1866 self.process_batch_from_right(right_batch)
1867 }
1868
1869 fn process_batch_after_right_end(
1870 &mut self,
1871 left_batch: &RecordBatch,
1872 ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
1873 self.process_batch_from_left(left_batch)
1874 }
1875
1876 fn process_batches_before_finalization(
1877 &mut self,
1878 ) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
1879 let left_result = build_side_determined_results(
1881 &self.left,
1882 &self.schema,
1883 self.left.input_buffer.num_rows(),
1884 self.right.input_buffer.schema(),
1885 self.join_type,
1886 &self.column_indices,
1887 )?;
1888 let right_result = build_side_determined_results(
1890 &self.right,
1891 &self.schema,
1892 self.right.input_buffer.num_rows(),
1893 self.left.input_buffer.schema(),
1894 self.join_type,
1895 &self.column_indices,
1896 )?;
1897
1898 let result = combine_two_batches(&self.schema, left_result, right_result)?;
1900
1901 if result.is_some() {
1903 return Ok(StatefulStreamResult::Ready(result));
1904 }
1905 Ok(StatefulStreamResult::Continue)
1906 }
1907
1908 fn right_stream(&mut self) -> &mut SendableRecordBatchStream {
1909 &mut self.right_stream
1910 }
1911
1912 fn left_stream(&mut self) -> &mut SendableRecordBatchStream {
1913 &mut self.left_stream
1914 }
1915
1916 fn set_state(&mut self, state: SHJStreamState) {
1917 self.state = state;
1918 }
1919
1920 fn state(&mut self) -> SHJStreamState {
1921 self.state.clone()
1922 }
1923
1924 fn size(&self) -> usize {
1925 let mut size = 0;
1926 size += size_of_val(&self.schema);
1927 size += size_of_val(&self.filter);
1928 size += size_of_val(&self.join_type);
1929 size += self.left.size();
1930 size += self.right.size();
1931 size += size_of_val(&self.column_indices);
1932 size += self.graph.as_ref().map(|g| g.size()).unwrap_or(0);
1933 size += size_of_val(&self.left_sorted_filter_expr);
1934 size += size_of_val(&self.right_sorted_filter_expr);
1935 size += size_of_val(&self.random_state);
1936 size += size_of_val(&self.null_equality);
1937 size += size_of_val(&self.metrics);
1938 size
1939 }
1940
1941 fn perform_join_for_given_side(
1949 &mut self,
1950 probe_batch: &RecordBatch,
1951 probe_side: JoinSide,
1952 ) -> Result<Option<RecordBatch>> {
1953 let (
1954 probe_hash_joiner,
1955 build_hash_joiner,
1956 probe_side_sorted_filter_expr,
1957 build_side_sorted_filter_expr,
1958 probe_side_metrics,
1959 ) = if probe_side.eq(&JoinSide::Left) {
1960 (
1961 &mut self.left,
1962 &mut self.right,
1963 &mut self.left_sorted_filter_expr,
1964 &mut self.right_sorted_filter_expr,
1965 &mut self.metrics.left,
1966 )
1967 } else {
1968 (
1969 &mut self.right,
1970 &mut self.left,
1971 &mut self.right_sorted_filter_expr,
1972 &mut self.left_sorted_filter_expr,
1973 &mut self.metrics.right,
1974 )
1975 };
1976 probe_side_metrics.input_batches.add(1);
1978 probe_side_metrics.input_rows.add(probe_batch.num_rows());
1979 probe_hash_joiner.update_internal_state(
1981 probe_batch,
1982 &self.random_state,
1983 self.null_equality,
1984 )?;
1985 let equal_result = join_with_probe_batch(
1987 build_hash_joiner,
1988 probe_hash_joiner,
1989 &self.schema,
1990 self.join_type,
1991 self.filter.as_ref(),
1992 probe_batch,
1993 &self.column_indices,
1994 &self.random_state,
1995 self.null_equality,
1996 )?;
1997 probe_hash_joiner.offset += probe_batch.num_rows();
1999
2000 let anti_result = if let (
2001 Some(build_side_sorted_filter_expr),
2002 Some(probe_side_sorted_filter_expr),
2003 Some(graph),
2004 ) = (
2005 build_side_sorted_filter_expr.as_mut(),
2006 probe_side_sorted_filter_expr.as_mut(),
2007 self.graph.as_mut(),
2008 ) {
2009 calculate_filter_expr_intervals(
2011 &build_hash_joiner.input_buffer,
2012 build_side_sorted_filter_expr,
2013 probe_batch,
2014 probe_side_sorted_filter_expr,
2015 )?;
2016 let prune_length = build_hash_joiner
2017 .calculate_prune_length_with_probe_batch(
2018 build_side_sorted_filter_expr,
2019 probe_side_sorted_filter_expr,
2020 graph,
2021 )?;
2022 let result = build_side_determined_results(
2023 build_hash_joiner,
2024 &self.schema,
2025 prune_length,
2026 probe_batch.schema(),
2027 self.join_type,
2028 &self.column_indices,
2029 )?;
2030 build_hash_joiner.prune_internal_state(prune_length)?;
2031 result
2032 } else {
2033 None
2034 };
2035
2036 let result = combine_two_batches(&self.schema, equal_result, anti_result)?;
2038 let capacity = self.size();
2039 self.metrics.stream_memory_usage.set(capacity);
2040 self.reservation.try_resize(capacity)?;
2041 Ok(result)
2042 }
2043}
2044
2045#[derive(Clone, Debug)]
2053pub enum SHJStreamState {
2054 PullRight,
2056
2057 PullLeft,
2059
2060 RightExhausted,
2062
2063 LeftExhausted,
2065
2066 BothExhausted { final_result: bool },
2071}
2072
2073#[cfg(test)]
2074mod tests {
2075 use std::collections::HashMap;
2076 use std::sync::{LazyLock, Mutex};
2077
2078 use super::*;
2079 use crate::joins::test_utils::{
2080 build_sides_record_batches, compare_batches, complicated_filter,
2081 create_memory_table, join_expr_tests_fixture_f64, join_expr_tests_fixture_i32,
2082 join_expr_tests_fixture_temporal, partitioned_hash_join_with_filter,
2083 partitioned_sym_join_with_filter, split_record_batches,
2084 };
2085
2086 use arrow::compute::SortOptions;
2087 use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit};
2088 use datafusion_common::ScalarValue;
2089 use datafusion_execution::config::SessionConfig;
2090 use datafusion_expr::Operator;
2091 use datafusion_physical_expr::expressions::{Column, binary, col, lit};
2092 use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
2093
2094 use rstest::*;
2095
2096 const TABLE_SIZE: i32 = 30;
2097
2098 type TableKey = (i32, i32, usize); type TableValue = (Vec<RecordBatch>, Vec<RecordBatch>); static TABLE_CACHE: LazyLock<Mutex<HashMap<TableKey, TableValue>>> =
2103 LazyLock::new(|| Mutex::new(HashMap::new()));
2104
2105 fn get_or_create_table(
2106 cardinality: (i32, i32),
2107 batch_size: usize,
2108 ) -> Result<TableValue> {
2109 {
2110 let cache = TABLE_CACHE.lock().unwrap();
2111 if let Some(table) = cache.get(&(cardinality.0, cardinality.1, batch_size)) {
2112 return Ok(table.clone());
2113 }
2114 }
2115
2116 let (left_batch, right_batch) =
2118 build_sides_record_batches(TABLE_SIZE, cardinality)?;
2119
2120 let (left_partition, right_partition) = (
2121 split_record_batches(&left_batch, batch_size)?,
2122 split_record_batches(&right_batch, batch_size)?,
2123 );
2124
2125 let mut cache = TABLE_CACHE.lock().unwrap();
2127
2128 cache.insert(
2130 (cardinality.0, cardinality.1, batch_size),
2131 (left_partition.clone(), right_partition.clone()),
2132 );
2133
2134 Ok((left_partition, right_partition))
2135 }
2136
2137 pub async fn experiment(
2138 left: Arc<dyn ExecutionPlan>,
2139 right: Arc<dyn ExecutionPlan>,
2140 filter: Option<JoinFilter>,
2141 join_type: JoinType,
2142 on: JoinOn,
2143 task_ctx: Arc<TaskContext>,
2144 ) -> Result<()> {
2145 let first_batches = partitioned_sym_join_with_filter(
2146 Arc::clone(&left),
2147 Arc::clone(&right),
2148 on.clone(),
2149 filter.clone(),
2150 &join_type,
2151 NullEquality::NullEqualsNothing,
2152 Arc::clone(&task_ctx),
2153 )
2154 .await?;
2155 let second_batches = partitioned_hash_join_with_filter(
2156 left,
2157 right,
2158 on,
2159 filter,
2160 &join_type,
2161 NullEquality::NullEqualsNothing,
2162 task_ctx,
2163 )
2164 .await?;
2165 compare_batches(&first_batches, &second_batches);
2166 Ok(())
2167 }
2168
2169 #[rstest]
2170 #[tokio::test(flavor = "multi_thread")]
2171 async fn complex_join_all_one_ascending_numeric(
2172 #[values(
2173 JoinType::Inner,
2174 JoinType::Left,
2175 JoinType::Right,
2176 JoinType::RightSemi,
2177 JoinType::LeftSemi,
2178 JoinType::LeftAnti,
2179 JoinType::LeftMark,
2180 JoinType::RightAnti,
2181 JoinType::RightMark,
2182 JoinType::Full
2183 )]
2184 join_type: JoinType,
2185 #[values(
2186 (4, 5),
2187 (12, 17),
2188 )]
2189 cardinality: (i32, i32),
2190 ) -> Result<()> {
2191 let task_ctx = Arc::new(TaskContext::default());
2193
2194 let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2195
2196 let left_schema = &left_partition[0].schema();
2197 let right_schema = &right_partition[0].schema();
2198
2199 let left_sorted = [PhysicalSortExpr {
2200 expr: binary(
2201 col("la1", left_schema)?,
2202 Operator::Plus,
2203 col("la2", left_schema)?,
2204 left_schema,
2205 )?,
2206 options: SortOptions::default(),
2207 }]
2208 .into();
2209 let right_sorted = [PhysicalSortExpr {
2210 expr: col("ra1", right_schema)?,
2211 options: SortOptions::default(),
2212 }]
2213 .into();
2214 let (left, right) = create_memory_table(
2215 left_partition,
2216 right_partition,
2217 vec![left_sorted],
2218 vec![right_sorted],
2219 )?;
2220
2221 let on = vec![(
2222 binary(
2223 col("lc1", left_schema)?,
2224 Operator::Plus,
2225 lit(ScalarValue::Int32(Some(1))),
2226 left_schema,
2227 )?,
2228 Arc::new(Column::new_with_schema("rc1", right_schema)?) as _,
2229 )];
2230
2231 let intermediate_schema = Schema::new(vec![
2232 Field::new("0", DataType::Int32, true),
2233 Field::new("1", DataType::Int32, true),
2234 Field::new("2", DataType::Int32, true),
2235 ]);
2236 let filter_expr = complicated_filter(&intermediate_schema)?;
2237 let column_indices = vec![
2238 ColumnIndex {
2239 index: left_schema.index_of("la1")?,
2240 side: JoinSide::Left,
2241 },
2242 ColumnIndex {
2243 index: left_schema.index_of("la2")?,
2244 side: JoinSide::Left,
2245 },
2246 ColumnIndex {
2247 index: right_schema.index_of("ra1")?,
2248 side: JoinSide::Right,
2249 },
2250 ];
2251 let filter =
2252 JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2253
2254 experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2255 Ok(())
2256 }
2257
2258 #[rstest]
2259 #[tokio::test(flavor = "multi_thread")]
2260 async fn join_all_one_ascending_numeric(
2261 #[values(
2262 JoinType::Inner,
2263 JoinType::Left,
2264 JoinType::Right,
2265 JoinType::RightSemi,
2266 JoinType::LeftSemi,
2267 JoinType::LeftAnti,
2268 JoinType::LeftMark,
2269 JoinType::RightAnti,
2270 JoinType::RightMark,
2271 JoinType::Full
2272 )]
2273 join_type: JoinType,
2274 #[values(0, 1, 2, 3, 4, 5)] case_expr: usize,
2275 ) -> Result<()> {
2276 let task_ctx = Arc::new(TaskContext::default());
2277 let (left_partition, right_partition) = get_or_create_table((4, 5), 8)?;
2278
2279 let left_schema = &left_partition[0].schema();
2280 let right_schema = &right_partition[0].schema();
2281
2282 let left_sorted = [PhysicalSortExpr {
2283 expr: col("la1", left_schema)?,
2284 options: SortOptions::default(),
2285 }]
2286 .into();
2287 let right_sorted = [PhysicalSortExpr {
2288 expr: col("ra1", right_schema)?,
2289 options: SortOptions::default(),
2290 }]
2291 .into();
2292 let (left, right) = create_memory_table(
2293 left_partition,
2294 right_partition,
2295 vec![left_sorted],
2296 vec![right_sorted],
2297 )?;
2298
2299 let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2300
2301 let intermediate_schema = Schema::new(vec![
2302 Field::new("left", DataType::Int32, true),
2303 Field::new("right", DataType::Int32, true),
2304 ]);
2305 let filter_expr = join_expr_tests_fixture_i32(
2306 case_expr,
2307 col("left", &intermediate_schema)?,
2308 col("right", &intermediate_schema)?,
2309 );
2310 let column_indices = vec![
2311 ColumnIndex {
2312 index: 0,
2313 side: JoinSide::Left,
2314 },
2315 ColumnIndex {
2316 index: 0,
2317 side: JoinSide::Right,
2318 },
2319 ];
2320 let filter =
2321 JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2322
2323 experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2324 Ok(())
2325 }
2326
2327 #[rstest]
2328 #[tokio::test(flavor = "multi_thread")]
2329 async fn join_without_sort_information(
2330 #[values(
2331 JoinType::Inner,
2332 JoinType::Left,
2333 JoinType::Right,
2334 JoinType::RightSemi,
2335 JoinType::LeftSemi,
2336 JoinType::LeftAnti,
2337 JoinType::LeftMark,
2338 JoinType::RightAnti,
2339 JoinType::RightMark,
2340 JoinType::Full
2341 )]
2342 join_type: JoinType,
2343 #[values(0, 1, 2, 3, 4, 5)] case_expr: usize,
2344 ) -> Result<()> {
2345 let task_ctx = Arc::new(TaskContext::default());
2346 let (left_partition, right_partition) = get_or_create_table((4, 5), 8)?;
2347
2348 let left_schema = &left_partition[0].schema();
2349 let right_schema = &right_partition[0].schema();
2350 let (left, right) =
2351 create_memory_table(left_partition, right_partition, vec![], vec![])?;
2352
2353 let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2354
2355 let intermediate_schema = Schema::new(vec![
2356 Field::new("left", DataType::Int32, true),
2357 Field::new("right", DataType::Int32, true),
2358 ]);
2359 let filter_expr = join_expr_tests_fixture_i32(
2360 case_expr,
2361 col("left", &intermediate_schema)?,
2362 col("right", &intermediate_schema)?,
2363 );
2364 let column_indices = vec![
2365 ColumnIndex {
2366 index: 5,
2367 side: JoinSide::Left,
2368 },
2369 ColumnIndex {
2370 index: 5,
2371 side: JoinSide::Right,
2372 },
2373 ];
2374 let filter =
2375 JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2376
2377 experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2378 Ok(())
2379 }
2380
2381 #[rstest]
2382 #[tokio::test(flavor = "multi_thread")]
2383 async fn join_without_filter(
2384 #[values(
2385 JoinType::Inner,
2386 JoinType::Left,
2387 JoinType::Right,
2388 JoinType::RightSemi,
2389 JoinType::LeftSemi,
2390 JoinType::LeftAnti,
2391 JoinType::LeftMark,
2392 JoinType::RightAnti,
2393 JoinType::RightMark,
2394 JoinType::Full
2395 )]
2396 join_type: JoinType,
2397 ) -> Result<()> {
2398 let task_ctx = Arc::new(TaskContext::default());
2399 let (left_partition, right_partition) = get_or_create_table((11, 21), 8)?;
2400 let left_schema = &left_partition[0].schema();
2401 let right_schema = &right_partition[0].schema();
2402 let (left, right) =
2403 create_memory_table(left_partition, right_partition, vec![], vec![])?;
2404
2405 let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2406 experiment(left, right, None, join_type, on, task_ctx).await?;
2407 Ok(())
2408 }
2409
2410 #[rstest]
2411 #[tokio::test(flavor = "multi_thread")]
2412 async fn join_all_one_descending_numeric_particular(
2413 #[values(
2414 JoinType::Inner,
2415 JoinType::Left,
2416 JoinType::Right,
2417 JoinType::RightSemi,
2418 JoinType::LeftSemi,
2419 JoinType::LeftAnti,
2420 JoinType::LeftMark,
2421 JoinType::RightAnti,
2422 JoinType::RightMark,
2423 JoinType::Full
2424 )]
2425 join_type: JoinType,
2426 #[values(0, 1, 2, 3, 4, 5)] case_expr: usize,
2427 ) -> Result<()> {
2428 let task_ctx = Arc::new(TaskContext::default());
2429 let (left_partition, right_partition) = get_or_create_table((11, 21), 8)?;
2430 let left_schema = &left_partition[0].schema();
2431 let right_schema = &right_partition[0].schema();
2432 let left_sorted = [PhysicalSortExpr {
2433 expr: col("la1_des", left_schema)?,
2434 options: SortOptions {
2435 descending: true,
2436 nulls_first: true,
2437 },
2438 }]
2439 .into();
2440 let right_sorted = [PhysicalSortExpr {
2441 expr: col("ra1_des", right_schema)?,
2442 options: SortOptions {
2443 descending: true,
2444 nulls_first: true,
2445 },
2446 }]
2447 .into();
2448 let (left, right) = create_memory_table(
2449 left_partition,
2450 right_partition,
2451 vec![left_sorted],
2452 vec![right_sorted],
2453 )?;
2454
2455 let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2456
2457 let intermediate_schema = Schema::new(vec![
2458 Field::new("left", DataType::Int32, true),
2459 Field::new("right", DataType::Int32, true),
2460 ]);
2461 let filter_expr = join_expr_tests_fixture_i32(
2462 case_expr,
2463 col("left", &intermediate_schema)?,
2464 col("right", &intermediate_schema)?,
2465 );
2466 let column_indices = vec![
2467 ColumnIndex {
2468 index: 5,
2469 side: JoinSide::Left,
2470 },
2471 ColumnIndex {
2472 index: 5,
2473 side: JoinSide::Right,
2474 },
2475 ];
2476 let filter =
2477 JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2478
2479 experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2480 Ok(())
2481 }
2482
2483 #[tokio::test(flavor = "multi_thread")]
2484 async fn build_null_columns_first() -> Result<()> {
2485 let join_type = JoinType::Full;
2486 let case_expr = 1;
2487 let session_config = SessionConfig::new().with_repartition_joins(false);
2488 let task_ctx = TaskContext::default().with_session_config(session_config);
2489 let task_ctx = Arc::new(task_ctx);
2490 let (left_partition, right_partition) = get_or_create_table((10, 11), 8)?;
2491 let left_schema = &left_partition[0].schema();
2492 let right_schema = &right_partition[0].schema();
2493 let left_sorted = [PhysicalSortExpr {
2494 expr: col("l_asc_null_first", left_schema)?,
2495 options: SortOptions {
2496 descending: false,
2497 nulls_first: true,
2498 },
2499 }]
2500 .into();
2501 let right_sorted = [PhysicalSortExpr {
2502 expr: col("r_asc_null_first", right_schema)?,
2503 options: SortOptions {
2504 descending: false,
2505 nulls_first: true,
2506 },
2507 }]
2508 .into();
2509 let (left, right) = create_memory_table(
2510 left_partition,
2511 right_partition,
2512 vec![left_sorted],
2513 vec![right_sorted],
2514 )?;
2515
2516 let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2517
2518 let intermediate_schema = Schema::new(vec![
2519 Field::new("left", DataType::Int32, true),
2520 Field::new("right", DataType::Int32, true),
2521 ]);
2522 let filter_expr = join_expr_tests_fixture_i32(
2523 case_expr,
2524 col("left", &intermediate_schema)?,
2525 col("right", &intermediate_schema)?,
2526 );
2527 let column_indices = vec![
2528 ColumnIndex {
2529 index: 6,
2530 side: JoinSide::Left,
2531 },
2532 ColumnIndex {
2533 index: 6,
2534 side: JoinSide::Right,
2535 },
2536 ];
2537 let filter =
2538 JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2539 experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2540 Ok(())
2541 }
2542
2543 #[tokio::test(flavor = "multi_thread")]
2544 async fn build_null_columns_last() -> Result<()> {
2545 let join_type = JoinType::Full;
2546 let case_expr = 1;
2547 let session_config = SessionConfig::new().with_repartition_joins(false);
2548 let task_ctx = TaskContext::default().with_session_config(session_config);
2549 let task_ctx = Arc::new(task_ctx);
2550 let (left_partition, right_partition) = get_or_create_table((10, 11), 8)?;
2551
2552 let left_schema = &left_partition[0].schema();
2553 let right_schema = &right_partition[0].schema();
2554 let left_sorted = [PhysicalSortExpr {
2555 expr: col("l_asc_null_last", left_schema)?,
2556 options: SortOptions {
2557 descending: false,
2558 nulls_first: false,
2559 },
2560 }]
2561 .into();
2562 let right_sorted = [PhysicalSortExpr {
2563 expr: col("r_asc_null_last", right_schema)?,
2564 options: SortOptions {
2565 descending: false,
2566 nulls_first: false,
2567 },
2568 }]
2569 .into();
2570 let (left, right) = create_memory_table(
2571 left_partition,
2572 right_partition,
2573 vec![left_sorted],
2574 vec![right_sorted],
2575 )?;
2576
2577 let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2578
2579 let intermediate_schema = Schema::new(vec![
2580 Field::new("left", DataType::Int32, true),
2581 Field::new("right", DataType::Int32, true),
2582 ]);
2583 let filter_expr = join_expr_tests_fixture_i32(
2584 case_expr,
2585 col("left", &intermediate_schema)?,
2586 col("right", &intermediate_schema)?,
2587 );
2588 let column_indices = vec![
2589 ColumnIndex {
2590 index: 7,
2591 side: JoinSide::Left,
2592 },
2593 ColumnIndex {
2594 index: 7,
2595 side: JoinSide::Right,
2596 },
2597 ];
2598 let filter =
2599 JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2600
2601 experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2602 Ok(())
2603 }
2604
2605 #[tokio::test(flavor = "multi_thread")]
2606 async fn build_null_columns_first_descending() -> Result<()> {
2607 let join_type = JoinType::Full;
2608 let cardinality = (10, 11);
2609 let case_expr = 1;
2610 let session_config = SessionConfig::new().with_repartition_joins(false);
2611 let task_ctx = TaskContext::default().with_session_config(session_config);
2612 let task_ctx = Arc::new(task_ctx);
2613 let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2614
2615 let left_schema = &left_partition[0].schema();
2616 let right_schema = &right_partition[0].schema();
2617 let left_sorted = [PhysicalSortExpr {
2618 expr: col("l_desc_null_first", left_schema)?,
2619 options: SortOptions {
2620 descending: true,
2621 nulls_first: true,
2622 },
2623 }]
2624 .into();
2625 let right_sorted = [PhysicalSortExpr {
2626 expr: col("r_desc_null_first", right_schema)?,
2627 options: SortOptions {
2628 descending: true,
2629 nulls_first: true,
2630 },
2631 }]
2632 .into();
2633 let (left, right) = create_memory_table(
2634 left_partition,
2635 right_partition,
2636 vec![left_sorted],
2637 vec![right_sorted],
2638 )?;
2639
2640 let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2641
2642 let intermediate_schema = Schema::new(vec![
2643 Field::new("left", DataType::Int32, true),
2644 Field::new("right", DataType::Int32, true),
2645 ]);
2646 let filter_expr = join_expr_tests_fixture_i32(
2647 case_expr,
2648 col("left", &intermediate_schema)?,
2649 col("right", &intermediate_schema)?,
2650 );
2651 let column_indices = vec![
2652 ColumnIndex {
2653 index: 8,
2654 side: JoinSide::Left,
2655 },
2656 ColumnIndex {
2657 index: 8,
2658 side: JoinSide::Right,
2659 },
2660 ];
2661 let filter =
2662 JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2663
2664 experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2665 Ok(())
2666 }
2667
2668 #[tokio::test(flavor = "multi_thread")]
2669 async fn complex_join_all_one_ascending_numeric_missing_stat() -> Result<()> {
2670 let cardinality = (3, 4);
2671 let join_type = JoinType::Full;
2672
2673 let session_config = SessionConfig::new().with_repartition_joins(false);
2675 let task_ctx = TaskContext::default().with_session_config(session_config);
2676 let task_ctx = Arc::new(task_ctx);
2677 let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2678
2679 let left_schema = &left_partition[0].schema();
2680 let right_schema = &right_partition[0].schema();
2681 let left_sorted = [PhysicalSortExpr {
2682 expr: col("la1", left_schema)?,
2683 options: SortOptions::default(),
2684 }]
2685 .into();
2686 let right_sorted = [PhysicalSortExpr {
2687 expr: col("ra1", right_schema)?,
2688 options: SortOptions::default(),
2689 }]
2690 .into();
2691 let (left, right) = create_memory_table(
2692 left_partition,
2693 right_partition,
2694 vec![left_sorted],
2695 vec![right_sorted],
2696 )?;
2697
2698 let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2699
2700 let intermediate_schema = Schema::new(vec![
2701 Field::new("0", DataType::Int32, true),
2702 Field::new("1", DataType::Int32, true),
2703 Field::new("2", DataType::Int32, true),
2704 ]);
2705 let filter_expr = complicated_filter(&intermediate_schema)?;
2706 let column_indices = vec![
2707 ColumnIndex {
2708 index: 0,
2709 side: JoinSide::Left,
2710 },
2711 ColumnIndex {
2712 index: 4,
2713 side: JoinSide::Left,
2714 },
2715 ColumnIndex {
2716 index: 0,
2717 side: JoinSide::Right,
2718 },
2719 ];
2720 let filter =
2721 JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2722
2723 experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2724 Ok(())
2725 }
2726
2727 #[tokio::test(flavor = "multi_thread")]
2728 async fn complex_join_all_one_ascending_equivalence() -> Result<()> {
2729 let cardinality = (3, 4);
2730 let join_type = JoinType::Full;
2731
2732 let config = SessionConfig::new().with_repartition_joins(false);
2734 let task_ctx = Arc::new(TaskContext::default().with_session_config(config));
2737 let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2738 let left_schema = &left_partition[0].schema();
2739 let right_schema = &right_partition[0].schema();
2740 let left_sorted = vec![
2741 [PhysicalSortExpr {
2742 expr: col("la1", left_schema)?,
2743 options: SortOptions::default(),
2744 }]
2745 .into(),
2746 [PhysicalSortExpr {
2747 expr: col("la2", left_schema)?,
2748 options: SortOptions::default(),
2749 }]
2750 .into(),
2751 ];
2752
2753 let right_sorted = [PhysicalSortExpr {
2754 expr: col("ra1", right_schema)?,
2755 options: SortOptions::default(),
2756 }]
2757 .into();
2758
2759 let (left, right) = create_memory_table(
2760 left_partition,
2761 right_partition,
2762 left_sorted,
2763 vec![right_sorted],
2764 )?;
2765
2766 let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2767
2768 let intermediate_schema = Schema::new(vec![
2769 Field::new("0", DataType::Int32, true),
2770 Field::new("1", DataType::Int32, true),
2771 Field::new("2", DataType::Int32, true),
2772 ]);
2773 let filter_expr = complicated_filter(&intermediate_schema)?;
2774 let column_indices = vec![
2775 ColumnIndex {
2776 index: 0,
2777 side: JoinSide::Left,
2778 },
2779 ColumnIndex {
2780 index: 4,
2781 side: JoinSide::Left,
2782 },
2783 ColumnIndex {
2784 index: 0,
2785 side: JoinSide::Right,
2786 },
2787 ];
2788 let filter =
2789 JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2790
2791 experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2792 Ok(())
2793 }
2794
2795 #[rstest]
2796 #[tokio::test(flavor = "multi_thread")]
2797 async fn testing_with_temporal_columns(
2798 #[values(
2799 JoinType::Inner,
2800 JoinType::Left,
2801 JoinType::Right,
2802 JoinType::RightSemi,
2803 JoinType::LeftSemi,
2804 JoinType::LeftAnti,
2805 JoinType::LeftMark,
2806 JoinType::RightAnti,
2807 JoinType::RightMark,
2808 JoinType::Full
2809 )]
2810 join_type: JoinType,
2811 #[values(
2812 (4, 5),
2813 (12, 17),
2814 )]
2815 cardinality: (i32, i32),
2816 #[values(0, 1, 2)] case_expr: usize,
2817 ) -> Result<()> {
2818 let session_config = SessionConfig::new().with_repartition_joins(false);
2819 let task_ctx = TaskContext::default().with_session_config(session_config);
2820 let task_ctx = Arc::new(task_ctx);
2821 let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2822
2823 let left_schema = &left_partition[0].schema();
2824 let right_schema = &right_partition[0].schema();
2825 let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2826 let left_sorted = [PhysicalSortExpr {
2827 expr: col("lt1", left_schema)?,
2828 options: SortOptions {
2829 descending: false,
2830 nulls_first: true,
2831 },
2832 }]
2833 .into();
2834 let right_sorted = [PhysicalSortExpr {
2835 expr: col("rt1", right_schema)?,
2836 options: SortOptions {
2837 descending: false,
2838 nulls_first: true,
2839 },
2840 }]
2841 .into();
2842 let (left, right) = create_memory_table(
2843 left_partition,
2844 right_partition,
2845 vec![left_sorted],
2846 vec![right_sorted],
2847 )?;
2848 let intermediate_schema = Schema::new(vec![
2849 Field::new(
2850 "left",
2851 DataType::Timestamp(TimeUnit::Millisecond, None),
2852 false,
2853 ),
2854 Field::new(
2855 "right",
2856 DataType::Timestamp(TimeUnit::Millisecond, None),
2857 false,
2858 ),
2859 ]);
2860 let filter_expr = join_expr_tests_fixture_temporal(
2861 case_expr,
2862 col("left", &intermediate_schema)?,
2863 col("right", &intermediate_schema)?,
2864 &intermediate_schema,
2865 )?;
2866 let column_indices = vec![
2867 ColumnIndex {
2868 index: 3,
2869 side: JoinSide::Left,
2870 },
2871 ColumnIndex {
2872 index: 3,
2873 side: JoinSide::Right,
2874 },
2875 ];
2876 let filter =
2877 JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2878 experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2879 Ok(())
2880 }
2881
2882 #[rstest]
2883 #[tokio::test(flavor = "multi_thread")]
2884 async fn test_with_interval_columns(
2885 #[values(
2886 JoinType::Inner,
2887 JoinType::Left,
2888 JoinType::Right,
2889 JoinType::RightSemi,
2890 JoinType::LeftSemi,
2891 JoinType::LeftAnti,
2892 JoinType::LeftMark,
2893 JoinType::RightAnti,
2894 JoinType::RightMark,
2895 JoinType::Full
2896 )]
2897 join_type: JoinType,
2898 #[values(
2899 (4, 5),
2900 (12, 17),
2901 )]
2902 cardinality: (i32, i32),
2903 ) -> Result<()> {
2904 let session_config = SessionConfig::new().with_repartition_joins(false);
2905 let task_ctx = TaskContext::default().with_session_config(session_config);
2906 let task_ctx = Arc::new(task_ctx);
2907 let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2908
2909 let left_schema = &left_partition[0].schema();
2910 let right_schema = &right_partition[0].schema();
2911 let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
2912 let left_sorted = [PhysicalSortExpr {
2913 expr: col("li1", left_schema)?,
2914 options: SortOptions {
2915 descending: false,
2916 nulls_first: true,
2917 },
2918 }]
2919 .into();
2920 let right_sorted = [PhysicalSortExpr {
2921 expr: col("ri1", right_schema)?,
2922 options: SortOptions {
2923 descending: false,
2924 nulls_first: true,
2925 },
2926 }]
2927 .into();
2928 let (left, right) = create_memory_table(
2929 left_partition,
2930 right_partition,
2931 vec![left_sorted],
2932 vec![right_sorted],
2933 )?;
2934 let intermediate_schema = Schema::new(vec![
2935 Field::new("left", DataType::Interval(IntervalUnit::DayTime), false),
2936 Field::new("right", DataType::Interval(IntervalUnit::DayTime), false),
2937 ]);
2938 let filter_expr = join_expr_tests_fixture_temporal(
2939 0,
2940 col("left", &intermediate_schema)?,
2941 col("right", &intermediate_schema)?,
2942 &intermediate_schema,
2943 )?;
2944 let column_indices = vec![
2945 ColumnIndex {
2946 index: 9,
2947 side: JoinSide::Left,
2948 },
2949 ColumnIndex {
2950 index: 9,
2951 side: JoinSide::Right,
2952 },
2953 ];
2954 let filter =
2955 JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
2956 experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
2957
2958 Ok(())
2959 }
2960
2961 #[rstest]
2962 #[tokio::test(flavor = "multi_thread")]
2963 async fn testing_ascending_float_pruning(
2964 #[values(
2965 JoinType::Inner,
2966 JoinType::Left,
2967 JoinType::Right,
2968 JoinType::RightSemi,
2969 JoinType::LeftSemi,
2970 JoinType::LeftAnti,
2971 JoinType::LeftMark,
2972 JoinType::RightAnti,
2973 JoinType::RightMark,
2974 JoinType::Full
2975 )]
2976 join_type: JoinType,
2977 #[values(
2978 (4, 5),
2979 (12, 17),
2980 )]
2981 cardinality: (i32, i32),
2982 #[values(0, 1, 2, 3, 4, 5)] case_expr: usize,
2983 ) -> Result<()> {
2984 let session_config = SessionConfig::new().with_repartition_joins(false);
2985 let task_ctx = TaskContext::default().with_session_config(session_config);
2986 let task_ctx = Arc::new(task_ctx);
2987 let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?;
2988
2989 let left_schema = &left_partition[0].schema();
2990 let right_schema = &right_partition[0].schema();
2991 let left_sorted = [PhysicalSortExpr {
2992 expr: col("l_float", left_schema)?,
2993 options: SortOptions::default(),
2994 }]
2995 .into();
2996 let right_sorted = [PhysicalSortExpr {
2997 expr: col("r_float", right_schema)?,
2998 options: SortOptions::default(),
2999 }]
3000 .into();
3001 let (left, right) = create_memory_table(
3002 left_partition,
3003 right_partition,
3004 vec![left_sorted],
3005 vec![right_sorted],
3006 )?;
3007
3008 let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)];
3009
3010 let intermediate_schema = Schema::new(vec![
3011 Field::new("left", DataType::Float64, true),
3012 Field::new("right", DataType::Float64, true),
3013 ]);
3014 let filter_expr = join_expr_tests_fixture_f64(
3015 case_expr,
3016 col("left", &intermediate_schema)?,
3017 col("right", &intermediate_schema)?,
3018 );
3019 let column_indices = vec![
3020 ColumnIndex {
3021 index: 10, side: JoinSide::Left,
3023 },
3024 ColumnIndex {
3025 index: 10, side: JoinSide::Right,
3027 },
3028 ];
3029 let filter =
3030 JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema));
3031
3032 experiment(left, right, Some(filter), join_type, on, task_ctx).await?;
3033 Ok(())
3034 }
3035}