2010YOUY01 commented on code in PR #24291:
URL: https://github.com/apache/datafusion/pull/24291#discussion_r3772272486
##########
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:
should we remove this `enforce_batch_size_in_joins`?
This should be introduced due to some legacy implementation tries to
materialize `total_left_buffer_len * right_batch_len` at once, now it only
materialize `one_left_row * right_batch_len` in each step, so the underlying
restriction is already followed
##########
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:
It seems we can put `fetch_probe_batch` inside `process_right_batch` to make
it even simpler
##########
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:
This slicing step seems not necessary. The `result` has length the same as
`right_batch`, and the upstream operator should already ensure it's less than
`batch_size` config (this is a global convention for operators)
##########
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:
I found this test and the next one is quite hard to figure out their test
goals, are they necessary? If they're auto generated we could remove them, else
it would be great to better explain their test goals.
##########
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:
```suggestion
// Cross join should follow the global output-size convention: all
output batches
// should have fewer than `batch_size` rows, as specified in the config.
async fn test_cross_join_output_length() -> Result<()> {
```
--
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]
