kosiew commented on code in PR #25372:
URL: https://github.com/apache/datafusion/pull/25372#discussion_r4266640292
##########
datafusion/physical-plan/src/sorts/stream.rs:
##########
@@ -587,4 +635,211 @@ mod tests {
Ok(())
}
+
+ /// Drives a `RowCursorStream` over several partitions and reports the pool
+ /// reservation, so the tests below can watch it move.
+ struct RowCursorHarness {
+ stream: RowCursorStream,
+ pool: Arc<dyn MemoryPool>,
+ schema: SchemaRef,
+ rows: usize,
+ str_width: usize,
+ }
+
+ /// Every batch has the same row count and string width, so each converts
to
+ /// a `Rows` buffer of the same size regardless of `seq`.
+ fn make_batch(
+ schema: &SchemaRef,
+ seq: usize,
+ rows: usize,
+ str_width: usize,
+ ) -> RecordBatch {
+ let base = (seq * rows) as i32;
+ let a = Int32Array::from_iter_values((0..rows as i32).map(|i| base +
i));
+ let b = StringArray::from_iter_values((0..rows).map(|_|
"x".repeat(str_width)));
+ RecordBatch::try_new(Arc::clone(schema), vec![Arc::new(a),
Arc::new(b)]).unwrap()
+ }
+
+ impl RowCursorHarness {
+ /// `batches_per_partition` batches of `rows` rows each, per partition.
+ /// The sort key spans two columns so this takes the `Rows` path rather
+ /// than the specialized single-column `FieldCursorStream`.
+ fn new(
+ partitions: usize,
+ batches_per_partition: usize,
+ rows: usize,
+ str_width: usize,
+ ) -> Result<Self> {
+ let schema: SchemaRef = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, false),
+ Field::new("b", DataType::Utf8, false),
+ ]));
+
+ let streams: Vec<SendableRecordBatchStream> = (0..partitions)
+ .map(|_| {
+ let batches: Vec<RecordBatch> = (0..batches_per_partition)
+ .map(|seq| make_batch(&schema, seq, rows, str_width))
+ .collect();
+ Box::pin(
+ MemoryStream::try_new(batches, Arc::clone(&schema),
None)
+ .unwrap(),
+ ) as SendableRecordBatchStream
+ })
+ .collect();
+
+ let expressions = LexOrdering::new(vec![
+ PhysicalSortExpr::new_default(col("a", &schema)?),
+ PhysicalSortExpr::new_default(col("b", &schema)?),
+ ])
+ .unwrap();
+
+ // Generous limit: this measures what is reserved, not what is
refused.
+ let pool: Arc<dyn MemoryPool> =
+ Arc::new(GreedyMemoryPool::new(256 * 1024 * 1024));
+ let reservation = MemoryConsumer::new("test").register(&pool);
+ let stream =
+ RowCursorStream::try_new(&schema, &expressions, streams,
reservation)?;
+
+ Ok(Self {
+ stream,
+ pool,
+ schema,
+ rows,
+ str_width,
+ })
+ }
+
+ fn reserved(&self) -> usize {
+ self.pool.reserved()
+ }
+
+ /// Converts one batch with a standalone converter over the same sort
+ /// key, returning `(converter size, size of one retained Rows
buffer)`.
+ /// These are what the stream should reserve for its converter and for
+ /// each partition's cached buffer.
+ fn expected_sizes(&self) -> Result<(usize, usize)> {
+ let converter = RowConverter::new(vec![
+ SortField::new(DataType::Int32),
+ SortField::new(DataType::Utf8),
+ ])?;
+ let batch = make_batch(&self.schema, 0, self.rows, self.str_width);
+ let mut rows = converter.empty_rows(0, 0);
+ converter.append(&mut rows, batch.columns())?;
+ Ok((converter.size(), rows.size()))
+ }
+
+ /// Polls `stream_idx`, handing back the cursor so the caller controls
when
+ /// it is dropped. `None` once the partition is finished.
+ fn poll_cursor(&mut self, stream_idx: usize) ->
Result<Option<RowValues>> {
+ let waker = futures::task::noop_waker();
+ let mut cx = Context::from_waker(&waker);
+ match self.stream.poll_next(&mut cx, stream_idx) {
+ Poll::Ready(Some(Ok((values, _batch)))) => Ok(Some(values)),
+ Poll::Ready(Some(Err(e))) => Err(e),
+ Poll::Ready(None) => Ok(None),
+ Poll::Pending => unreachable!("MemoryStream is never pending"),
+ }
+ }
+
+ /// Polls `stream_idx` and immediately drops the returned cursor,
which is
+ /// what the merge does once a cursor is exhausted. Returns `false`
once
+ /// the partition is finished.
+ fn poll_and_drop_cursor(&mut self, stream_idx: usize) -> Result<bool> {
+ let waker = futures::task::noop_waker();
+ let mut cx = Context::from_waker(&waker);
+ match self.stream.poll_next(&mut cx, stream_idx) {
+ Poll::Ready(Some(Ok((values, _batch)))) => {
+ drop(values);
+ Ok(true)
+ }
+ Poll::Ready(Some(Err(e))) => Err(e),
+ Poll::Ready(None) => Ok(false),
+ Poll::Pending => unreachable!("MemoryStream is never pending"),
+ }
+ }
+ }
+
+ // `RowCursorStream` keeps one `Rows` buffer per partition alive between
+ // polls so the allocation can be reused. That buffer is real memory, and
it
+ // outlives the cursor handed to the consumer: the merge drops a cursor as
+ // soon as it is exhausted, but the buffer behind it stays cached.
+ //
+ // So dropping the cursor must not change what the pool reports. If it
does,
+ // the bytes still held by the cache have gone off the books and the pool
is
+ // under-counting a live allocation.
+ #[test]
+ fn dropping_a_cursor_does_not_unaccount_its_retained_buffer() ->
Result<()> {
+ let mut harness = RowCursorHarness::new(1, 4, 512, 64)?;
+ let (converter_size, rows_size) = harness.expected_sizes()?;
+ assert_eq!(harness.reserved(), 0);
+
+ let cursor = harness.poll_cursor(0)?.expect("first batch");
+ assert_eq!(harness.reserved(), converter_size + rows_size);
+
+ // The cached `Rows` buffer outlives the cursor, so dropping the cursor
+ // must not release its bytes.
+ drop(cursor);
+ assert_eq!(harness.reserved(), converter_size + rows_size);
+
+ Ok(())
+ }
+
+ // The retained bytes are per partition, so what the pool reports has to
+ // scale with partition count rather than staying flat at the size of
+ // whichever batch happens to be in flight.
+ #[test]
+ fn retained_row_buffer_reservation_scales_with_partitions() -> Result<()> {
+ for partitions in [2, 16] {
+ let mut harness = RowCursorHarness::new(partitions, 4, 512, 64)?;
+ let (converter_size, rows_size) = harness.expected_sizes()?;
+ for stream_idx in 0..partitions {
+ assert!(harness.poll_and_drop_cursor(stream_idx)?);
+ }
+ // One shared converter, plus one cached buffer per partition.
+ assert_eq!(
+ harness.reserved(),
+ converter_size + partitions * rows_size,
+ "{partitions} partitions"
+ );
+ }
+
+ Ok(())
+ }
+
+ // The flip side: a partition that will never be polled again must hand its
+ // buffer back. Otherwise the reservation only ever grows and a long merge
+ // holds every partition's high-water mark until the whole stream drops.
+ #[test]
+ fn exhausted_partitions_release_their_retained_buffers() -> Result<()> {
Review Comment:
Could we add a small, large, small batch sequence while another partition
keeps its buffer cached? This would verify that the reservation grows with the
actual retained `Rows::size()`, stays at the high-water mark after the smaller
batch, and shrinks correctly on exhaustion without affecting the other
partition.
--
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]
