(datafusion) branch main updated: Reuse spill_record_batch_by_size function (#12389)

[co...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

comphead 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 09f7592bd6 Reuse spill_record_batch_by_size function (#12389)
09f7592bd6 is described below

commit 09f7592bd6ebbeacaa1515d2f2d3a79dcc5a9178
Author: 张林伟 <le...@gmail.com>
AuthorDate: Tue Sep 10 00:34:12 2024 +0800

    Reuse spill_record_batch_by_size function (#12389)
---
 datafusion/physical-plan/src/aggregates/mod.rs      |  8 ++++----
 datafusion/physical-plan/src/aggregates/row_hash.rs | 21 +++++++--------------
 datafusion/physical-plan/src/spill.rs               |  1 -
 3 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index ad0e43631a..c3bc7b042e 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -1136,12 +1136,12 @@ fn evaluate_optional(
 /// Evaluate a group by expression against a `RecordBatch`
 ///
 /// Arguments:
-/// `group_by`: the expression to evaluate
-/// `batch`: the `RecordBatch` to evaluate against
+/// - `group_by`: the expression to evaluate
+/// - `batch`: the `RecordBatch` to evaluate against
 ///
 /// Returns: A Vec of Vecs of Array of results
-/// The outer Vect appears to be for grouping sets
-/// The inner Vect contains the results per expression
+/// The outer Vec appears to be for grouping sets
+/// The inner Vec contains the results per expression
 /// The inner-inner Array contains the results per row
 pub(crate) fn evaluate_group_by(
     group_by: &PhysicalGroupBy,
diff --git a/datafusion/physical-plan/src/aggregates/row_hash.rs b/datafusion/physical-plan/src/aggregates/row_hash.rs
index fd2d26d9b4..0332131d4b 100644
--- a/datafusion/physical-plan/src/aggregates/row_hash.rs
+++ b/datafusion/physical-plan/src/aggregates/row_hash.rs
@@ -27,11 +27,10 @@ use crate::aggregates::{
     evaluate_group_by, evaluate_many, evaluate_optional, group_schema, AggregateMode,
     PhysicalGroupBy,
 };
-use crate::common::IPCWriter;
 use crate::metrics::{BaselineMetrics, MetricBuilder, RecordOutput};
 use crate::sorts::sort::sort_batch;
 use crate::sorts::streaming_merge;
-use crate::spill::read_spill_as_stream;
+use crate::spill::{read_spill_as_stream, spill_record_batch_by_size};
 use crate::stream::RecordBatchStreamAdapter;
 use crate::{aggregates, metrics, ExecutionPlan, PhysicalExpr};
 use crate::{RecordBatchStream, SendableRecordBatchStream};
@@ -905,19 +904,13 @@ impl GroupedHashAggregateStream {
         let emit = self.emit(EmitTo::All, true)?;
         let sorted = sort_batch(&emit, &self.spill_state.spill_expr, None)?;
         let spillfile = self.runtime.disk_manager.create_tmp_file("HashAggSpill")?;
-        let mut writer = IPCWriter::new(spillfile.path(), &emit.schema())?;
         // TODO: slice large `sorted` and write to multiple files in parallel
-        let mut offset = 0;
-        let total_rows = sorted.num_rows();
-
-        while offset < total_rows {
-            let length = std::cmp::min(total_rows - offset, self.batch_size);
-            let batch = sorted.slice(offset, length);
-            offset += batch.num_rows();
-            writer.write(&batch)?;
-        }
-
-        writer.finish()?;
+        spill_record_batch_by_size(
+            &sorted,
+            spillfile.path().into(),
+            sorted.schema(),
+            self.batch_size,
+        )?;
         self.spill_state.spills.push(spillfile);
         Ok(())
     }
diff --git a/datafusion/physical-plan/src/spill.rs b/datafusion/physical-plan/src/spill.rs
index 21ca58fa0a..de85a7c6f0 100644
--- a/datafusion/physical-plan/src/spill.rs
+++ b/datafusion/physical-plan/src/spill.rs
@@ -88,7 +88,6 @@ fn read_spill(sender: Sender<Result<RecordBatch>>, path: &Path) -> Result<()> {
 
 /// Spill the `RecordBatch` to disk as smaller batches
 /// split by `batch_size_rows`
-/// Return `total_rows` what is spilled
 pub fn spill_record_batch_by_size(
     batch: &RecordBatch,
     path: PathBuf,


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org