saadtajwar commented on code in PR #24291:
URL: https://github.com/apache/datafusion/pull/24291#discussion_r3777175791
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -1054,6 +991,196 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn test_join_enforce_batch_size_splits_output() -> Result<()> {
Review Comment:
Removed this test as we removed the batch sizing!
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -361,29 +364,25 @@ impl ExecutionPlan for CrossJoinExec {
))
})?;
- if enforce_batch_size_in_joins {
- Ok(Box::pin(CrossJoinStream {
- schema: Arc::clone(&self.schema),
- left_fut,
- right: stream,
- left_index: 0,
- join_metrics,
- state: CrossJoinStreamState::WaitBuildSide,
- left_data: RecordBatch::new_empty(self.left().schema()),
- batch_transformer: BatchSplitter::new(batch_size),
- }))
- } else {
- Ok(Box::pin(CrossJoinStream {
- schema: Arc::clone(&self.schema),
- left_fut,
- right: stream,
- left_index: 0,
- join_metrics,
- state: CrossJoinStreamState::WaitBuildSide,
- left_data: RecordBatch::new_empty(self.left().schema()),
- batch_transformer: NoopBatchTransformer::new(),
- }))
- }
+ let mut state = CrossJoinStream {
+ schema: Arc::clone(&self.schema),
+ left_fut,
+ right: stream,
+ join_metrics,
+ left_data: RecordBatch::new_empty(self.left().schema()),
+ batch_size: enforce_batch_size_in_joins.then_some(batch_size),
Review Comment:
Ooh I see, yes that makes sense to me!
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -641,129 +611,96 @@ fn build_batch(
.map_err(Into::into)
}
-#[async_trait]
-impl<T: BatchTransformer + Unpin + Send> Stream for CrossJoinStream<T> {
- type Item = Result<RecordBatch>;
-
- fn poll_next(
- mut self: std::pin::Pin<&mut Self>,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Self::Item>> {
- self.poll_next_impl(cx)
- }
-}
-
-impl<T: BatchTransformer> CrossJoinStream<T> {
- /// Separate implementation function that unpins the [`CrossJoinStream`] so
- /// that partial borrows work correctly
- fn poll_next_impl(
+impl CrossJoinStream {
+ // Collect the left (build) side, then continue processing the right side
against it until we have no more rows on the right
+ async fn join(
&mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
- loop {
- return match self.state {
- CrossJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
- }
- CrossJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
- }
- CrossJoinStreamState::BuildBatches(_) => {
- let poll = handle_state!(self.build_batches());
- self.join_metrics.baseline.record_poll(poll)
- }
- };
+ emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
+ ) -> Result<()> {
+ if !self.collect_build_side().await? {
+ return Ok(());
+ }
+
+ while let Some(right_batch) = self.fetch_probe_batch().await? {
+ self.process_right_batch(&right_batch, emitter).await?
Review Comment:
Ooh yeah - done!
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -1054,6 +991,196 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn test_join_enforce_batch_size_splits_output() -> Result<()> {
+ let mut config = SessionConfig::new().with_batch_size(2);
+ config.options_mut().execution.enforce_batch_size_in_joins = true;
+ let task_ctx =
Arc::new(TaskContext::default().with_session_config(config));
+
+ let left = build_table_scan_i32(
+ ("a1", &vec![1, 2, 3]),
+ ("b1", &vec![4, 5, 6]),
+ ("c1", &vec![7, 8, 9]),
+ );
+ let right = build_table_scan_i32(
+ ("a2", &vec![10, 11, 12, 13, 14]),
+ ("b2", &vec![15, 16, 17, 18, 19]),
+ ("c2", &vec![20, 21, 22, 23, 24]),
+ );
+
+ let (_, batches, _) = join_collect(left, right, task_ctx).await?;
+
+ let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+ assert_eq!(total_rows, 15);
+ assert!(batches.iter().all(|b| b.num_rows() <= 2));
+
+ Ok(())
+ }
+
+ fn delayed_stream(
+ batches: Vec<RecordBatch>,
+ delay: Duration,
+ ) -> SendableRecordBatchStream {
+ let schema = batches[0].schema();
+ Box::pin(RecordBatchStreamAdapter::new(
+ schema,
+ futures::stream::iter(batches.into_iter().map(Ok)).then(
+ move |item| async move {
+ tokio::time::sleep(delay).await;
+ item
+ },
+ ),
+ ))
+ }
+
+ fn probe_batches() -> Vec<RecordBatch> {
+ vec![
+ build_table_i32(
+ ("a2", &vec![10, 11]),
+ ("b2", &vec![12, 13]),
+ ("c2", &vec![14, 15]),
+ ),
+ build_table_i32(
+ ("a2", &vec![20, 21]),
+ ("b2", &vec![22, 23]),
+ ("c2", &vec![24, 25]),
+ ),
+ build_table_i32(
+ ("a2", &vec![30, 31]),
+ ("b2", &vec![32, 33]),
+ ("c2", &vec![34, 35]),
+ ),
+ ]
+ }
+
+ fn cross_join_stream(
+ left_batch: RecordBatch,
+ right: SendableRecordBatchStream,
+ batch_size: Option<usize>,
+ ) -> Result<(SendableRecordBatchStream, ExecutionPlanMetricsSet)> {
+ let metrics = ExecutionPlanMetricsSet::new();
+ let join_metrics = BuildProbeJoinMetrics::new(0, &metrics);
+
+ let runtime = RuntimeEnvBuilder::new().build_arc()?;
+ let reservation =
MemoryConsumer::new("test").register(&runtime.memory_pool);
+
+ let left_schema = left_batch.schema();
+ let mut fields: Vec<_> =
left_schema.fields().iter().cloned().collect();
+ fields.extend(right.schema().fields().iter().cloned());
+ let schema = Arc::new(Schema::new(fields));
+
+ let left_data = JoinLeftData {
+ merged_batch: left_batch,
+ _reservation: reservation,
+ };
+ let left_fut = OnceFut::new(async move { Ok(left_data) });
+
+ let mut state = CrossJoinStream {
+ schema: Arc::clone(&schema),
+ left_fut,
+ right,
+ join_metrics,
+ left_data: RecordBatch::new_empty(left_schema),
+ batch_size,
+ };
+ let baseline_metrics = state.join_metrics.baseline.clone();
+ let stream =
+ async_try_stream(
+ move |mut emitter| async move { state.join(&mut emitter).await
},
+ );
+ let observed = ObservedStream::new(
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream)),
+ baseline_metrics,
+ None,
+ );
+ Ok((Box::pin(observed), metrics))
+ }
+
+ fn elapsed_compute_of(metrics: &ExecutionPlanMetricsSet) -> Duration {
+
Duration::from_nanos(metrics.clone_inner().elapsed_compute().unwrap_or(0) as
u64)
+ }
+
+ async fn check_elapsed_compute_excluded<F, Fut>(mut run: F) -> Result<()>
+ where
+ F: FnMut(Duration) -> Fut,
+ Fut: Future<Output = Result<(Duration, Duration)>>,
+ {
+ let mut delay = Duration::from_millis(50);
+ for attempt in 0..3 {
+ let (elapsed_compute, wall) = run(delay).await?;
+ if elapsed_compute < delay {
+ return Ok(());
+ }
+ assert!(
+ attempt < 2,
+ "elapsed_compute ({elapsed_compute:?}) should be well below
the \
+ injected delay ({delay:?}); wall {wall:?}"
+ );
+ delay *= 4;
+ }
+ unreachable!()
+ }
+
+ #[tokio::test]
+ async fn elapsed_compute_excludes_probe_input_wait() -> Result<()> {
Review Comment:
Removed these :)
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -641,129 +611,96 @@ fn build_batch(
.map_err(Into::into)
}
-#[async_trait]
-impl<T: BatchTransformer + Unpin + Send> Stream for CrossJoinStream<T> {
- type Item = Result<RecordBatch>;
-
- fn poll_next(
- mut self: std::pin::Pin<&mut Self>,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Self::Item>> {
- self.poll_next_impl(cx)
- }
-}
-
-impl<T: BatchTransformer> CrossJoinStream<T> {
- /// Separate implementation function that unpins the [`CrossJoinStream`] so
- /// that partial borrows work correctly
- fn poll_next_impl(
+impl CrossJoinStream {
+ // Collect the left (build) side, then continue processing the right side
against it until we have no more rows on the right
+ async fn join(
&mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
- loop {
- return match self.state {
- CrossJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
- }
- CrossJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
- }
- CrossJoinStreamState::BuildBatches(_) => {
- let poll = handle_state!(self.build_batches());
- self.join_metrics.baseline.record_poll(poll)
- }
- };
+ emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
+ ) -> Result<()> {
+ if !self.collect_build_side().await? {
+ return Ok(());
+ }
+
+ while let Some(right_batch) = self.fetch_probe_batch().await? {
+ self.process_right_batch(&right_batch, emitter).await?
}
+
+ Ok(())
}
- /// Collects build (left) side of the join into the state. In case of an
empty build batch,
- /// the execution terminates. Otherwise, the state is updated to fetch
probe (right) batch.
- fn collect_build_side(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- let build_timer = self.join_metrics.build_time.timer();
- let left_data = match ready!(self.left_fut.get(cx)) {
- Ok(left_data) => left_data,
- Err(e) => return Poll::Ready(Err(e)),
- };
- build_timer.done();
-
- let left_data = left_data.merged_batch.clone();
- let result = if left_data.num_rows() == 0 {
- StatefulStreamResult::Ready(None)
- } else {
- self.left_data = left_data;
- self.state = CrossJoinStreamState::FetchProbeBatch;
- StatefulStreamResult::Continue
- };
- Poll::Ready(Ok(result))
+ /// Collects build (left) side of the join into the state. In case of an
empty build batch, the execution terminates.
+ /// Returns true if build side was loaded and non-empty
+ async fn collect_build_side(&mut self) -> Result<bool> {
+ let left_data = poll_fn(|cx| {
+ self.left_fut
+ .get(cx)
+ .map(|res| res.map(|data| data.merged_batch.clone()))
+ })
+ .await?;
+
+ let is_empty = left_data.num_rows().is_zero();
+ self.left_data = left_data;
+ Ok(!is_empty)
}
- /// Fetches the probe (right) batch, updates the metrics, and save the
batch in the state.
- /// Then, the state is updated to build result batches.
- fn fetch_probe_batch(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- self.left_index = 0;
- let right_data = match ready!(self.right.poll_next_unpin(cx)) {
+ /// Fetches the probe (right) batch, updates the metrics, and returns the
batch
+ async fn fetch_probe_batch(&mut self) -> Result<Option<RecordBatch>> {
+ let right_data = match self.right.next().await {
Some(Ok(right_data)) => right_data,
- Some(Err(e)) => return Poll::Ready(Err(e)),
+ Some(Err(e)) => return Err(e),
None => {
// Release the right (probe) input pipeline's resources.
let right_schema = self.right.schema();
self.right =
Box::pin(EmptyRecordBatchStream::new(right_schema));
- return Poll::Ready(Ok(StatefulStreamResult::Ready(None)));
+ return Ok(None);
}
};
self.join_metrics.input_batches.add(1);
self.join_metrics.input_rows.add(right_data.num_rows());
- self.state = CrossJoinStreamState::BuildBatches(right_data);
- Poll::Ready(Ok(StatefulStreamResult::Continue))
+ Ok(Some(right_data))
}
- /// Joins the indexed row of left data with the current probe batch.
- /// If all the results are produced, the state is set to fetch new probe
batch.
- fn build_batches(&mut self) ->
Result<StatefulStreamResult<Option<RecordBatch>>> {
- let right_batch = self.state.try_as_record_batch()?;
- if self.left_index < self.left_data.num_rows() {
- match self.batch_transformer.next() {
- None => {
- let join_timer = self.join_metrics.join_time.timer();
- let result = build_batch(
- self.left_index,
- right_batch,
- &self.left_data,
- &self.schema,
- );
- join_timer.done();
-
- self.batch_transformer.set_batch(result?);
- }
- Some((batch, last)) => {
- if last {
- self.left_index += 1;
- }
-
- return Ok(StatefulStreamResult::Ready(Some(batch)));
+ /// Joins the left data with the current probe batch, using the emitter to
emit the resultant batches
+ async fn process_right_batch(
+ &mut self,
+ right_batch: &RecordBatch,
+ emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
+ ) -> Result<()> {
+ for left_index in 0..self.left_data.num_rows() {
+ let join_timer = self.join_metrics.join_time.timer();
+ let result =
+ build_batch(left_index, right_batch, &self.left_data,
&self.schema)?;
+ join_timer.done();
+
+ if let Some(batch_size) = self.batch_size {
Review Comment:
Removed the batch sizing from the first comment :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
