Copilot commented on code in PR #25565:
URL: https://github.com/apache/datafusion/pull/25565#discussion_r4093032619
##########
datafusion/physical-plan/src/sorts/builder.rs:
##########
@@ -272,13 +334,238 @@ impl BatchBuilder {
return Ok(None);
}
+ let Some((target_batch_bytes, mut estimated_bytes)) =
+ self.target_batch_bytes.and_then(|target| {
+ self.estimated_prefix_bytes(self.indices.len())
+ .map(|estimated| (target, estimated))
+ })
+ else {
+ let (rows_to_emit, columns) = retry_interleave(
+ self.indices.len(),
+ self.indices.len(),
+ |rows_to_emit|
self.try_interleave_columns(&self.indices[..rows_to_emit]),
+ )?;
+
+ return Ok(Some(self.finish_record_batch(rows_to_emit, columns)?));
+ };
+
+ let initial_rows_to_emit = self.indices.len();
+ let mut rows_to_emit =
+ if initial_rows_to_emit <= 1 || estimated_bytes <=
target_batch_bytes {
+ initial_rows_to_emit
+ } else {
+ self.largest_prefix_under_target(target_batch_bytes)
+ .unwrap_or(1)
+ };
+
+ if rows_to_emit != initial_rows_to_emit {
+ estimated_bytes = self
+ .estimated_prefix_bytes(rows_to_emit)
+ .expect("a smaller prefix has the same supported arrays");
+ }
+
+ loop {
+ match try_grow_reservation_to_at_least(
+ &mut self.output_construction_reservation,
+ estimated_bytes,
+ ) {
+ Ok(()) => break,
+ Err(_) if rows_to_emit > 1 => {
+ let failed_bytes = estimated_bytes;
+ rows_to_emit /= 2;
+ estimated_bytes = self
+ .estimated_prefix_bytes(rows_to_emit)
+ .expect("a smaller prefix has the same supported
arrays");
+ warn!(
+ "Could not reserve {failed_bytes} bytes for sort
output, retrying with {rows_to_emit} rows requiring {estimated_bytes} bytes"
+ );
+ }
+ Err(e) => return Err(e),
+ }
+ }
+
let (rows_to_emit, columns) =
- retry_interleave(self.indices.len(), self.indices.len(),
|rows_to_emit| {
+ match retry_interleave(rows_to_emit, rows_to_emit, |rows_to_emit| {
self.try_interleave_columns(&self.indices[..rows_to_emit])
- })?;
+ }) {
+ Ok(value) => value,
+ Err(e) => {
+ self.output_construction_reservation.free();
+ return Err(e);
+ }
+ };
Ok(Some(self.finish_record_batch(rows_to_emit, columns)?))
}
+
+ fn largest_prefix_under_target(&self, target_batch_bytes: usize) ->
Option<usize> {
+ if self.estimated_prefix_bytes(1)? > target_batch_bytes {
+ return Some(1);
+ }
+
+ let mut low = 1;
+ let mut high = self.indices.len();
+ while low < high {
+ let mid = low + (high - low).div_ceil(2);
+ if self.estimated_prefix_bytes(mid)? <= target_batch_bytes {
+ low = mid;
+ } else {
+ high = mid - 1;
+ }
+ }
+ Some(low)
+ }
+
+ fn estimated_prefix_bytes(&self, rows_to_emit: usize) -> Option<usize> {
+ let mut total = 0usize;
+ for column_idx in 0..self.schema.fields.len() {
+ total = total.checked_add(
+ self.column_prefix_memory_upper_bound(column_idx,
rows_to_emit)?,
+ )?;
+ }
+ Some(total)
+ }
+
+ fn column_prefix_memory_upper_bound(
+ &self,
+ column_idx: usize,
+ rows_to_emit: usize,
+ ) -> Option<usize> {
+ let data_type = self.schema.field(column_idx).data_type();
+
+ match data_type {
+ DataType::Null => Some(0),
+ DataType::Boolean => bitmap_buffer_bytes(rows_to_emit)?
+ .checked_add(self.validity_buffer_bytes(column_idx,
rows_to_emit)?),
+ DataType::Binary =>
self.byte_array_prefix_memory_upper_bound::<BinaryType>(
+ column_idx,
+ rows_to_emit,
+ ),
+ DataType::LargeBinary => self
+ .byte_array_prefix_memory_upper_bound::<LargeBinaryType>(
+ column_idx,
+ rows_to_emit,
+ ),
+ DataType::Utf8 =>
self.byte_array_prefix_memory_upper_bound::<Utf8Type>(
+ column_idx,
+ rows_to_emit,
+ ),
+ DataType::LargeUtf8 => self
+ .byte_array_prefix_memory_upper_bound::<LargeUtf8Type>(
+ column_idx,
+ rows_to_emit,
+ ),
+ DataType::FixedSizeBinary(width) => usize::try_from(*width)
+ .ok()
+ .and_then(|width| {
+ rows_to_emit
+ .checked_mul(width)
+ .and_then(aligned_buffer_bytes)
+ })
+ .and_then(|values| {
+ self.validity_buffer_bytes(column_idx, rows_to_emit)?
+ .checked_add(values)
+ }),
+ _ => fixed_width(data_type).and_then(|width| {
+ rows_to_emit.checked_mul(width).and_then(|values| {
+ self.validity_buffer_bytes(column_idx, rows_to_emit)?
+ .checked_add(values)
+ })
+ }),
+ }
+ }
+
+ /// Arrow allocates an output validity buffer when any input array for this
+ /// column contains nulls, even if the selected rows are all valid.
+ fn validity_buffer_bytes(
+ &self,
+ column_idx: usize,
+ rows_to_emit: usize,
+ ) -> Option<usize> {
+ if self
+ .batches
+ .iter()
+ .any(|(_, batch)| batch.column(column_idx).null_count() > 0)
+ {
+ bitmap_buffer_bytes(rows_to_emit)
+ } else {
+ Some(0)
+ }
+ }
+
+ fn byte_array_prefix_memory_upper_bound<T: ByteArrayType>(
+ &self,
+ column_idx: usize,
+ rows_to_emit: usize,
+ ) -> Option<usize> {
+ let mut values_len = 0usize;
+ for (batch_idx, row_idx) in &self.indices[..rows_to_emit] {
+ let array = self.batches[*batch_idx].1.column(column_idx);
+ let array = array.as_any().downcast_ref::<GenericByteArray<T>>()?;
+ values_len =
+
values_len.checked_add(array.value_length(*row_idx).as_usize())?;
+ }
+
+ self.validity_buffer_bytes(column_idx, rows_to_emit)?
+ .checked_add((rows_to_emit +
1).checked_mul(size_of::<T::Offset>())?)?
+ .checked_add(values_len)
Review Comment:
`values_len` is not an upper bound on the allocation made for an interleaved
byte-array values buffer: Arrow's mutable byte buffer can retain alignment
padding, just as the `FixedSizeBinary` branch accounts for above. This can
reserve the unaligned estimate, allocate the batch, and only then fail (or emit
a multi-row batch above the byte target) when `finish_record_batch` observes
the larger capacity. Round the values-buffer estimate before using it for
prefix selection and pre-allocation admission.
--
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]
