memory limited hash join (#5490)

datafusion/core/src/physical_plan/common.rs

@@ -39,8 +39,14 @@ use std::task::{Context, Poll};
 use tokio::sync::mpsc;
 use tokio::task::JoinHandle;
 
+/// [`MemoryReservation`] used across query execution streams
 pub(crate) type SharedMemoryReservation = Arc<Mutex<MemoryReservation>>;
 
+/// [`MemoryReservation`] used at query operator level
+/// `Option` wrapper allows to initialize empty reservation in operator constructor,
+/// and set it to actual reservation at stream level.
+pub(crate) type OperatorMemoryReservation = Arc<Mutex<Option<SharedMemoryReservation>>>;
+
 /// Stream of record batches
 pub struct SizedRecordBatchStream {
  schema: SchemaRef,

datafusion/core/src/physical_plan/joins/cross_join.rs

@@ -27,7 +27,7 @@ use arrow::record_batch::RecordBatch;
 
 use crate::execution::context::TaskContext;
 use crate::execution::memory_pool::MemoryConsumer;
-use crate::physical_plan::common::SharedMemoryReservation;
+use crate::physical_plan::common::{OperatorMemoryReservation, SharedMemoryReservation};
 use crate::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
 use crate::physical_plan::{
  coalesce_batches::concat_batches, coalesce_partitions::CoalescePartitionsExec,
@@ -60,6 +60,8 @@ pub struct CrossJoinExec {
  schema: SchemaRef,
  /// Build-side data
  left_fut: OnceAsync<JoinLeftData>,
+ /// Memory reservation for build-side data
+ reservation: OperatorMemoryReservation,
  /// Execution plan metrics
  metrics: ExecutionPlanMetricsSet,
 }
@@ -83,6 +85,7 @@ impl CrossJoinExec {
  right,
  schema,
  left_fut: Default::default(),
+ reservation: Default::default(),
  metrics: ExecutionPlanMetricsSet::default(),
  }
  }
@@ -221,17 +224,29 @@ impl ExecutionPlan for CrossJoinExec {
  let stream = self.right.execute(partition, context.clone())?;
 
  let join_metrics = BuildProbeJoinMetrics::new(partition, &self.metrics);
- let reservation = Arc::new(Mutex::new(
- MemoryConsumer::new(format!("CrossJoinStream[{partition}]"))
- .register(context.memory_pool()),
- ));
+
+ // Initialization of operator-level reservation
+ {
+ let mut reservation_lock = self.reservation.lock();
+ if reservation_lock.is_none() {
+ *reservation_lock = Some(Arc::new(Mutex::new(
+ MemoryConsumer::new("CrossJoinExec").register(context.memory_pool()),
+ )));
+ };
+ }
+
+ let reservation = self.reservation.lock().clone().ok_or_else(|| {
+ DataFusionError::Internal(
+ "Operator-level memory reservation is not initialized".to_string(),
+ )
+ })?;
 
  let left_fut = self.left_fut.once(|| {
  load_left_input(
  self.left.clone(),
  context,
  join_metrics.clone(),
- reservation.clone(),
+ reservation,
  )
  });
 
@@ -242,7 +257,6 @@ impl ExecutionPlan for CrossJoinExec {
  right_batch: Arc::new(parking_lot::Mutex::new(None)),
  left_index: 0,
  join_metrics,
- reservation,
  }))
  }
 
@@ -346,8 +360,6 @@ struct CrossJoinStream {
  right_batch: Arc<parking_lot::Mutex<Option<RecordBatch>>>,
  /// join execution metrics
  join_metrics: BuildProbeJoinMetrics,
- /// memory reservation
- reservation: SharedMemoryReservation,
 }
 
 impl RecordBatchStream for CrossJoinStream {
@@ -452,10 +464,7 @@ impl CrossJoinStream {
 
  Some(result)
  }
- other => {
- self.reservation.lock().free();
- other
- }
+ other => other,
  })
  }
 }
@@ -683,6 +692,7 @@ mod tests {
  err.to_string(),
  "External error: Resources exhausted: Failed to allocate additional"
  );
+ assert_contains!(err.to_string(), "CrossJoinExec");
 
  Ok(())
  }