This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 72c6491d25 Minor: Use upstream concat_batches (#11615)
72c6491d25 is described below
commit 72c6491d25fe253b8757028be77e1e6f5cd74c71
Author: Andrew Lamb <[email protected]>
AuthorDate: Wed Jul 24 08:36:49 2024 -0400
Minor: Use upstream concat_batches (#11615)
---
datafusion/physical-plan/src/coalesce_batches.rs | 38 ++++------------------
datafusion/physical-plan/src/joins/cross_join.rs | 9 +++--
.../physical-plan/src/joins/nested_loop_join.rs | 6 ++--
3 files changed, 14 insertions(+), 39 deletions(-)
diff --git a/datafusion/physical-plan/src/coalesce_batches.rs
b/datafusion/physical-plan/src/coalesce_batches.rs
index b9bdfcdee7..8cb25827ff 100644
--- a/datafusion/physical-plan/src/coalesce_batches.rs
+++ b/datafusion/physical-plan/src/coalesce_batches.rs
@@ -18,25 +18,23 @@
//! CoalesceBatchesExec combines small batches into larger batches for more
efficient use of
//! vectorized processing by upstream operators.
-use std::any::Any;
-use std::pin::Pin;
-use std::sync::Arc;
-use std::task::{Context, Poll};
-
use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use super::{DisplayAs, ExecutionPlanProperties, PlanProperties, Statistics};
use crate::{
DisplayFormatType, ExecutionPlan, RecordBatchStream,
SendableRecordBatchStream,
};
+use arrow::compute::concat_batches;
+use std::any::Any;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
use arrow::datatypes::SchemaRef;
-use arrow::error::Result as ArrowResult;
use arrow::record_batch::RecordBatch;
use datafusion_common::Result;
use datafusion_execution::TaskContext;
use futures::stream::{Stream, StreamExt};
-use log::trace;
/// CoalesceBatchesExec combines small batches into larger batches for more
efficient use of
/// vectorized processing by upstream operators.
@@ -229,11 +227,7 @@ impl CoalesceBatchesStream {
// check to see if we have enough batches yet
if self.buffered_rows >= self.target_batch_size {
// combine the batches and return
- let batch = concat_batches(
- &self.schema,
- &self.buffer,
- self.buffered_rows,
- )?;
+ let batch = concat_batches(&self.schema,
&self.buffer)?;
// reset buffer state
self.buffer.clear();
self.buffered_rows = 0;
@@ -250,11 +244,7 @@ impl CoalesceBatchesStream {
return Poll::Ready(None);
} else {
// combine the batches and return
- let batch = concat_batches(
- &self.schema,
- &self.buffer,
- self.buffered_rows,
- )?;
+ let batch = concat_batches(&self.schema,
&self.buffer)?;
// reset buffer state
self.buffer.clear();
self.buffered_rows = 0;
@@ -276,20 +266,6 @@ impl RecordBatchStream for CoalesceBatchesStream {
}
}
-/// Concatenates an array of `RecordBatch` into one batch
-pub fn concat_batches(
- schema: &SchemaRef,
- batches: &[RecordBatch],
- row_count: usize,
-) -> ArrowResult<RecordBatch> {
- trace!(
- "Combined {} batches containing {} rows",
- batches.len(),
- row_count
- );
- arrow::compute::concat_batches(schema, batches)
-}
-
#[cfg(test)]
mod tests {
use super::*;
diff --git a/datafusion/physical-plan/src/joins/cross_join.rs
b/datafusion/physical-plan/src/joins/cross_join.rs
index 8304ddc733..b1482a9699 100644
--- a/datafusion/physical-plan/src/joins/cross_join.rs
+++ b/datafusion/physical-plan/src/joins/cross_join.rs
@@ -18,13 +18,10 @@
//! Defines the cross join plan for loading the left side of the cross join
//! and producing batches in parallel for the right partitions
-use std::{any::Any, sync::Arc, task::Poll};
-
use super::utils::{
adjust_right_output_partitioning, BuildProbeJoinMetrics, OnceAsync,
OnceFut,
StatefulStreamResult,
};
-use crate::coalesce_batches::concat_batches;
use crate::coalesce_partitions::CoalescePartitionsExec;
use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet};
use crate::{
@@ -33,6 +30,8 @@ use crate::{
ExecutionPlanProperties, PlanProperties, RecordBatchStream,
SendableRecordBatchStream, Statistics,
};
+use arrow::compute::concat_batches;
+use std::{any::Any, sync::Arc, task::Poll};
use arrow::datatypes::{Fields, Schema, SchemaRef};
use arrow::record_batch::RecordBatch;
@@ -155,7 +154,7 @@ async fn load_left_input(
let stream = merge.execute(0, context)?;
// Load all batches and count the rows
- let (batches, num_rows, _, reservation) = stream
+ let (batches, _num_rows, _, reservation) = stream
.try_fold(
(Vec::new(), 0usize, metrics, reservation),
|mut acc, batch| async {
@@ -175,7 +174,7 @@ async fn load_left_input(
)
.await?;
- let merged_batch = concat_batches(&left_schema, &batches, num_rows)?;
+ let merged_batch = concat_batches(&left_schema, &batches)?;
Ok((merged_batch, reservation))
}
diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs
b/datafusion/physical-plan/src/joins/nested_loop_join.rs
index f8ca389808..eac135bfd0 100644
--- a/datafusion/physical-plan/src/joins/nested_loop_join.rs
+++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs
@@ -26,7 +26,6 @@ use std::sync::Arc;
use std::task::Poll;
use super::utils::{asymmetric_join_output_partitioning,
need_produce_result_in_final};
-use crate::coalesce_batches::concat_batches;
use crate::coalesce_partitions::CoalescePartitionsExec;
use crate::joins::utils::{
adjust_indices_by_join_type, apply_join_filter_to_indices,
build_batch_from_indices,
@@ -44,6 +43,7 @@ use crate::{
use arrow::array::{
BooleanBufferBuilder, UInt32Array, UInt32Builder, UInt64Array,
UInt64Builder,
};
+use arrow::compute::concat_batches;
use arrow::datatypes::{Schema, SchemaRef};
use arrow::record_batch::RecordBatch;
use arrow::util::bit_util;
@@ -364,7 +364,7 @@ async fn collect_left_input(
let stream = merge.execute(0, context)?;
// Load all batches and count the rows
- let (batches, num_rows, metrics, mut reservation) = stream
+ let (batches, _num_rows, metrics, mut reservation) = stream
.try_fold(
(Vec::new(), 0usize, join_metrics, reservation),
|mut acc, batch| async {
@@ -384,7 +384,7 @@ async fn collect_left_input(
)
.await?;
- let merged_batch = concat_batches(&schema, &batches, num_rows)?;
+ let merged_batch = concat_batches(&schema, &batches)?;
// Reserve memory for visited_left_side bitmap if required by join type
let visited_left_side = if with_visited_left_side {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
