2010YOUY01 commented on code in PR #25820:
URL: https://github.com/apache/datafusion/pull/25820#discussion_r4117802191
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -3763,6 +3899,60 @@ mod tests {
Ok(())
}
+ /// Zero-copy slices of one batch share its buffers, so buffering them
+ /// reserves those buffers once. Sorting the slices as separate runs must
+ /// keep the shared buffers reserved while any run still holds them.
+ #[tokio::test]
+ async fn test_sliced_runs_reserve_shared_buffers_once() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![Field::new("x",
DataType::Int64, false)]));
+ let parent = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(Int64Array::from_iter_values((0..4096).rev()))],
+ )?;
+ let slices: Vec<_> = (0..4).map(|i| parent.slice(i * 1024,
1024)).collect();
+
+ let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1024 *
1024));
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_pool(Arc::clone(&pool))
+ .build_arc()?;
+ let ordering: LexOrdering =
+ [PhysicalSortExpr::new_default(Arc::new(Column::new("x",
0)))].into();
+ let mut sorter = ExternalSorter::new(
+ 0,
+ Arc::clone(&schema),
+ ordering.clone(),
+ 1024,
+ 0,
+ 0, // Sort each slice as its own run, rather than concatenating
them.
+ SpillCompression::Uncompressed,
+ &ExecutionPlanMetricsSet::new(),
+ runtime,
+ )?;
+ for slice in &slices {
+ sorter.insert_batch(slice.clone()).await?;
+ }
+ let parent_bytes = get_record_batch_memory_size(&parent);
+ let sliced_bytes = slices
+ .iter()
+ .map(|slice| slice.get_sliced_size())
+ .sum::<Result<usize>>()?;
+ assert_eq!(pool.reserved(), parent_bytes + sliced_bytes);
+
+ let runs = std::mem::take(&mut sorter.in_mem_batches);
+ let mut streams = sorter.sort_run_streams(runs)?;
+ let last = streams.pop().unwrap();
+ for stream in streams {
+ stream.try_collect::<Vec<_>>().await?;
+ }
+ // The last run still holds the parent's buffers until it is sorted
+ assert_eq!(pool.reserved(), parent_bytes +
slices[3].get_sliced_size()?);
Review Comment:
Could we move this test to e2e tests, like running a sort query with memory
pool limit 10MB (which fails without PR)
Probably also good to just remove it, since there are already e2e coverage
in this PR.
This test is asserting implementation details, which is likely to get
refactored often
- Reference:
https://github.com/apache/datafusion/blob/main/docs/source/contributor-guide/testing.md#choosing-what-kind-of-test-to-write
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -865,7 +941,10 @@ impl ExternalSorter {
&mut self,
input: &RecordBatch,
) -> Result<()> {
- let size = get_reserved_bytes_for_record_batch(input)?;
+ let size = get_reserved_bytes_for_next_record_batch(
Review Comment:
follow-up idea: I find this part hard to reason about, it's like estimating
the memory delta given input batches might have shared buffer.
Maybe we can estimate it differently to simplify it, e.g. always estimate
the total memory usage, and use `try_resize()` instead of `try_grow()`
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -946,6 +1030,58 @@ pub(crate) fn get_reserved_bytes_for_record_batch(batch:
&RecordBatch) -> Result
})
}
+/// Estimate how much memory is needed to sort `batch` when it is buffered with
+/// the batches `counter` has already counted.
+///
+/// Like [`get_reserved_bytes_for_record_batch`], but only counts the buffers
of
+/// `batch` that `counter` has not seen. Buffered batches can share buffers,
for
+/// example zero-copy slices of one larger batch, and a shared buffer only
needs
+/// to be reserved once. Without a `counter`, every buffer of `batch` is
counted.
+fn get_reserved_bytes_for_next_record_batch(
+ batch: &RecordBatch,
+ counter: Option<&mut RecordBatchMemoryCounter>,
+) -> Result<usize> {
+ let Some(counter) = counter else {
+ return get_reserved_bytes_for_record_batch(batch);
+ };
+ let sliced_size = batch.get_sliced_size()?;
+ Ok(get_reserved_bytes_for_record_batch_size(
+ counter.count_batch(batch),
+ sliced_size,
+ ))
+}
+
+/// Returns whether the sorted copy of an array of `data_type` keeps
+/// referencing some of the array's buffers.
+///
+/// Sorting takes the selected rows into new buffers, except that `take` keeps
+/// the data buffers of view arrays and the values of dictionaries and list
+/// views. Every sorted run is charged for those buffers in full, so buffered
+/// batches that share them must be charged for them each, too.
+fn sorted_copy_keeps_buffers(data_type: &DataType) -> bool {
Review Comment:
nit and optional: we could remove it and always assume there are shared
buffers to keep it simple, unless there are proven slowdowns.
--
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]
