Dandandan commented on code in PR #16647:
URL: https://github.com/apache/datafusion/pull/16647#discussion_r2180844635
##########
datafusion/physical-plan/src/sorts/stream.rs:
##########
@@ -105,26 +110,53 @@ impl RowCursorStream {
})
.collect::<Result<Vec<_>>>()?;
- let streams = streams.into_iter().map(|s| s.fuse()).collect();
+ let streams: Vec<_> = streams.into_iter().map(|s| s.fuse()).collect();
let converter = RowConverter::new(sort_fields)?;
+ let mut rows = Vec::with_capacity(streams.len());
+ for _ in &streams {
+ // Initialize each stream with an empty Rows
+ rows.push([
+ Some(Arc::new(converter.empty_rows(0, 0))),
+ Some(Arc::new(converter.empty_rows(0, 0))),
+ ]);
+ }
Ok(Self {
converter,
reservation,
column_expressions: expressions.iter().map(|x|
Arc::clone(&x.expr)).collect(),
streams: FusedStreams(streams),
+ rows,
})
}
- fn convert_batch(&mut self, batch: &RecordBatch) -> Result<RowValues> {
+ fn convert_batch(
+ &mut self,
+ batch: &RecordBatch,
+ stream_idx: usize,
+ ) -> Result<RowValues> {
let cols = self
.column_expressions
.iter()
.map(|expr| expr.evaluate(batch)?.into_array(batch.num_rows()))
.collect::<Result<Vec<_>>>()?;
- let rows = self.converter.convert_columns(&cols)?;
+ // At this point, ownership should of this Rows should be unique
+ let mut rows =
Arc::try_unwrap(self.rows[stream_idx][1].take().unwrap())
+ .unwrap_or_else(|_| self.converter.empty_rows(0, 0));
+
+ rows.clear();
+
+ self.converter.append(&mut rows, &cols)?;
self.reservation.try_resize(self.converter.size())?;
+ let rows = Arc::new(rows);
+
+ self.rows[stream_idx][1] = Some(Arc::clone(&rows));
+
+ // swap the curent with the previous one, so that the next poll can
reuse the Rows from the previous poll
+ let [a, b] = &mut self.rows[stream_idx];
+ std::mem::swap(a, b);
Review Comment:
Yeah > cannot borrow `self.rows` as mutable more than once at a time
--
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: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
