Skip to content

Commit db51150

Browse files
ashdnazgashdnazg
committed
fix: overcounting of memory in first/last.
When aggregating first/last list over a column of lists, the first/last accumulators hold the necessary scalar value as is, which points to the list in the original input buffer. This results in two issues: 1) We prevent the deallocation of the input arrays which might be significantly larger than the single value we want to hold. 2) During aggreagtion with groups, many accumulators receive slices of the same input buffer, resulting in all held values pointing to this buffer. Then, when calculating the size of all accumulators we count the buffer multiple times, since each accumulator considers it to be part of its own allocation.
1 parent f12ba60 commit db51150

File tree

3 files changed

+165
-50
lines changed

3 files changed

+165
-50
lines changed

‎datafusion/common/src/scalar/mod.rs‎

Lines changed: 61 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -3435,49 +3435,80 @@ impl ScalarValue {
34353435
.sum::<usize>()
34363436
}
34373437

3438-
/// Performs a deep clone of the ScalarValue, creating new copies of all nested data structures.
3439-
/// This is different from the standard `clone()` which may share data through `Arc`.
3440-
/// Aggregation functions like `max` will cost a lot of memory if the data is not cloned.
3441-
pub fn force_clone(&self) -> Self {
3438+
/// Compacts the allocation referenced by `self` to the minimum, copying the data if
3439+
/// necessary.
3440+
///
3441+
/// This can be relevant when `self` is a list or contains a list as a nested value, as
3442+
/// a single list holds an Arc to its entire original array buffer.
3443+
pub fn compact(&mut self) {
34423444
match self {
3443-
// Complex types need deep clone of their contents
3444-
ScalarValue::List(array) => {
3445-
let array = copy_array_data(&array.to_data());
3446-
let new_array = ListArray::from(array);
3447-
ScalarValue::List(Arc::new(new_array))
3445+
ScalarValue::Null
3446+
| ScalarValue::Boolean(_)
3447+
| ScalarValue::Float16(_)
3448+
| ScalarValue::Float32(_)
3449+
| ScalarValue::Float64(_)
3450+
| ScalarValue::Decimal128(_, _, _)
3451+
| ScalarValue::Decimal256(_, _, _)
3452+
| ScalarValue::Int8(_)
3453+
| ScalarValue::Int16(_)
3454+
| ScalarValue::Int32(_)
3455+
| ScalarValue::Int64(_)
3456+
| ScalarValue::UInt8(_)
3457+
| ScalarValue::UInt16(_)
3458+
| ScalarValue::UInt32(_)
3459+
| ScalarValue::UInt64(_)
3460+
| ScalarValue::Date32(_)
3461+
| ScalarValue::Date64(_)
3462+
| ScalarValue::Time32Second(_)
3463+
| ScalarValue::Time32Millisecond(_)
3464+
| ScalarValue::Time64Microsecond(_)
3465+
| ScalarValue::Time64Nanosecond(_)
3466+
| ScalarValue::IntervalYearMonth(_)
3467+
| ScalarValue::IntervalDayTime(_)
3468+
| ScalarValue::IntervalMonthDayNano(_)
3469+
| ScalarValue::DurationSecond(_)
3470+
| ScalarValue::DurationMillisecond(_)
3471+
| ScalarValue::DurationMicrosecond(_)
3472+
| ScalarValue::DurationNanosecond(_)
3473+
| ScalarValue::Utf8(_)
3474+
| ScalarValue::LargeUtf8(_)
3475+
| ScalarValue::Utf8View(_)
3476+
| ScalarValue::TimestampSecond(_, _)
3477+
| ScalarValue::TimestampMillisecond(_, _)
3478+
| ScalarValue::TimestampMicrosecond(_, _)
3479+
| ScalarValue::TimestampNanosecond(_, _)
3480+
| ScalarValue::Binary(_)
3481+
| ScalarValue::FixedSizeBinary(_, _)
3482+
| ScalarValue::LargeBinary(_)
3483+
| ScalarValue::BinaryView(_) => (),
3484+
ScalarValue::FixedSizeList(arr) => {
3485+
let array = copy_array_data(&arr.to_data());
3486+
*Arc::make_mut(arr) = FixedSizeListArray::from(array);
34483487
}
3449-
ScalarValue::LargeList(array) => {
3450-
let array = copy_array_data(&array.to_data());
3451-
let new_array = LargeListArray::from(array);
3452-
ScalarValue::LargeList(Arc::new(new_array))
3488+
ScalarValue::List(arr) => {
3489+
let array = copy_array_data(&arr.to_data());
3490+
*Arc::make_mut(arr) = ListArray::from(array);
34533491
}
3454-
ScalarValue::FixedSizeList(arr) => {
3492+
ScalarValue::LargeList(arr) => {
34553493
let array = copy_array_data(&arr.to_data());
3456-
let new_array = FixedSizeListArray::from(array);
3457-
ScalarValue::FixedSizeList(Arc::new(new_array))
3494+
*Arc::make_mut(arr) = LargeListArray::from(array)
34583495
}
34593496
ScalarValue::Struct(arr) => {
34603497
let array = copy_array_data(&arr.to_data());
3461-
let new_array = StructArray::from(array);
3462-
ScalarValue::Struct(Arc::new(new_array))
3498+
*Arc::make_mut(arr) = StructArray::from(array);
34633499
}
34643500
ScalarValue::Map(arr) => {
34653501
let array = copy_array_data(&arr.to_data());
3466-
let new_array = MapArray::from(array);
3467-
ScalarValue::Map(Arc::new(new_array))
3468-
}
3469-
ScalarValue::Union(Some((type_id, value)), fields, mode) => {
3470-
let new_value = Box::new(value.force_clone());
3471-
ScalarValue::Union(Some((*type_id, new_value)), fields.clone(), *mode)
3502+
*Arc::make_mut(arr) = MapArray::from(array);
34723503
}
3473-
ScalarValue::Union(None, fields, mode) => {
3474-
ScalarValue::Union(None, fields.clone(), *mode)
3504+
ScalarValue::Union(val, _, _) => {
3505+
if let Some((_, value)) = val.as_mut() {
3506+
value.compact();
3507+
}
34753508
}
3476-
ScalarValue::Dictionary(key_type, value) => {
3477-
let new_value = Box::new(value.force_clone());
3478-
ScalarValue::Dictionary(key_type.clone(), new_value)
3509+
ScalarValue::Dictionary(_, value) => {
3510+
value.compact();
34793511
}
3480-
_ => self.clone(),
34813512
}
34823513
}
34833514
}

‎datafusion/functions-aggregate/src/first_last.rs‎

Lines changed: 90 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -752,7 +752,7 @@ where
752752

753753
fn size(&self) -> usize {
754754
self.vals.capacity() * size_of::<T::Native>()
755-
+ self.null_builder.capacity() / 8 // capacity is in bits, so convert to bytes
755+
+ self.null_builder.capacity() / 8 // capacity is in bits, so convert to bytes
756756
+ self.is_sets.capacity() / 8
757757
+ self.size_of_orderings
758758
+ self.min_of_each_group_buf.0.capacity() * size_of::<usize>()
@@ -827,9 +827,13 @@ impl FirstValueAccumulator {
827827
}
828828

829829
// Updates state with the values in the given row.
830-
fn update_with_new_row(&mut self, row: &[ScalarValue]) {
831-
self.first = row[0].clone();
832-
self.orderings = row[1..].to_vec();
830+
fn update_with_new_row(&mut self, mut row: Vec<ScalarValue>) {
831+
row.iter_mut().for_each(|s| {
832+
s.compact();
833+
});
834+
835+
self.first = row.remove(0);
836+
self.orderings = row;
833837
self.is_set = true;
834838
}
835839

@@ -888,7 +892,7 @@ impl Accumulator for FirstValueAccumulator {
888892
if !self.is_set {
889893
if let Some(first_idx) = self.get_first_idx(values)? {
890894
let row = get_row_at_idx(values, first_idx)?;
891-
self.update_with_new_row(&row);
895+
self.update_with_new_row(row);
892896
}
893897
} else if !self.requirement_satisfied {
894898
if let Some(first_idx) = self.get_first_idx(values)? {
@@ -901,7 +905,7 @@ impl Accumulator for FirstValueAccumulator {
901905
)?
902906
.is_gt()
903907
{
904-
self.update_with_new_row(&row);
908+
self.update_with_new_row(row);
905909
}
906910
}
907911
}
@@ -925,7 +929,7 @@ impl Accumulator for FirstValueAccumulator {
925929
let min = (0..filtered_states[0].len()).min_by(|&a, &b| comparator.compare(a, b));
926930

927931
if let Some(first_idx) = min {
928-
let first_row = get_row_at_idx(&filtered_states, first_idx)?;
932+
let mut first_row = get_row_at_idx(&filtered_states, first_idx)?;
929933
// When collecting orderings, we exclude the is_set flag from the state.
930934
let first_ordering = &first_row[1..is_set_idx];
931935
let sort_options = get_sort_options(self.ordering_req.as_ref());
@@ -936,7 +940,9 @@ impl Accumulator for FirstValueAccumulator {
936940
// Update with first value in the state. Note that we should exclude the
937941
// is_set flag from the state. Otherwise, we will end up with a state
938942
// containing two is_set flags.
939-
self.update_with_new_row(&first_row[0..is_set_idx]);
943+
assert!(is_set_idx <= first_row.len());
944+
first_row.resize(is_set_idx, ScalarValue::Null);
945+
self.update_with_new_row(first_row);
940946
}
941947
}
942948
Ok(())
@@ -1226,9 +1232,13 @@ impl LastValueAccumulator {
12261232
}
12271233

12281234
// Updates state with the values in the given row.
1229-
fn update_with_new_row(&mut self, row: &[ScalarValue]) {
1230-
self.last = row[0].clone();
1231-
self.orderings = row[1..].to_vec();
1235+
fn update_with_new_row(&mut self, mut row: Vec<ScalarValue>) {
1236+
row.iter_mut().for_each(|s| {
1237+
s.compact();
1238+
});
1239+
1240+
self.last = row.remove(0);
1241+
self.orderings = row;
12321242
self.is_set = true;
12331243
}
12341244

@@ -1289,7 +1299,7 @@ impl Accumulator for LastValueAccumulator {
12891299
if !self.is_set || self.requirement_satisfied {
12901300
if let Some(last_idx) = self.get_last_idx(values)? {
12911301
let row = get_row_at_idx(values, last_idx)?;
1292-
self.update_with_new_row(&row);
1302+
self.update_with_new_row(row);
12931303
}
12941304
} else if let Some(last_idx) = self.get_last_idx(values)? {
12951305
let row = get_row_at_idx(values, last_idx)?;
@@ -1302,7 +1312,7 @@ impl Accumulator for LastValueAccumulator {
13021312
)?
13031313
.is_lt()
13041314
{
1305-
self.update_with_new_row(&row);
1315+
self.update_with_new_row(row);
13061316
}
13071317
}
13081318

@@ -1326,7 +1336,7 @@ impl Accumulator for LastValueAccumulator {
13261336
let max = (0..filtered_states[0].len()).max_by(|&a, &b| comparator.compare(a, b));
13271337

13281338
if let Some(last_idx) = max {
1329-
let last_row = get_row_at_idx(&filtered_states, last_idx)?;
1339+
let mut last_row = get_row_at_idx(&filtered_states, last_idx)?;
13301340
// When collecting orderings, we exclude the is_set flag from the state.
13311341
let last_ordering = &last_row[1..is_set_idx];
13321342
let sort_options = get_sort_options(self.ordering_req.as_ref());
@@ -1339,7 +1349,9 @@ impl Accumulator for LastValueAccumulator {
13391349
// Update with last value in the state. Note that we should exclude the
13401350
// is_set flag from the state. Otherwise, we will end up with a state
13411351
// containing two is_set flags.
1342-
self.update_with_new_row(&last_row[0..is_set_idx]);
1352+
assert!(is_set_idx <= last_row.len());
1353+
last_row.resize(is_set_idx, ScalarValue::Null);
1354+
self.update_with_new_row(last_row);
13431355
}
13441356
}
13451357
Ok(())
@@ -1382,7 +1394,13 @@ fn convert_to_sort_cols(arrs: &[ArrayRef], sort_exprs: &LexOrdering) -> Vec<Sort
13821394

13831395
#[cfg(test)]
13841396
mod tests {
1385-
use arrow::{array::Int64Array, compute::SortOptions, datatypes::Schema};
1397+
use std::iter::repeat_with;
1398+
1399+
use arrow::{
1400+
array::{Int64Array, ListArray},
1401+
compute::SortOptions,
1402+
datatypes::Schema,
1403+
};
13861404
use datafusion_physical_expr::{expressions::col, PhysicalSortExpr};
13871405

13881406
use super::*;
@@ -1772,4 +1790,60 @@ mod tests {
17721790

17731791
Ok(())
17741792
}
1793+
1794+
#[test]
1795+
fn test_first_list_acc_size() -> Result<()> {
1796+
fn size_after_batch(values: &[ArrayRef]) -> Result<usize> {
1797+
let mut first_accumulator = FirstValueAccumulator::try_new(
1798+
&DataType::List(Arc::new(Field::new_list_field(DataType::Int64, false))),
1799+
&[],
1800+
LexOrdering::default(),
1801+
false,
1802+
)?;
1803+
1804+
first_accumulator.update_batch(values)?;
1805+
1806+
Ok(first_accumulator.size())
1807+
}
1808+
1809+
let batch1 = ListArray::from_iter_primitive::<Int32Type, _, _>(
1810+
repeat_with(|| Some(vec![Some(1)])).take(10000),
1811+
);
1812+
let batch2 =
1813+
ListArray::from_iter_primitive::<Int32Type, _, _>([Some(vec![Some(1)])]);
1814+
1815+
let size1 = size_after_batch(&[Arc::new(batch1)])?;
1816+
let size2 = size_after_batch(&[Arc::new(batch2)])?;
1817+
assert_eq!(size1, size2);
1818+
1819+
Ok(())
1820+
}
1821+
1822+
#[test]
1823+
fn test_last_list_acc_size() -> Result<()> {
1824+
fn size_after_batch(values: &[ArrayRef]) -> Result<usize> {
1825+
let mut last_accumulator = FirstValueAccumulator::try_new(
1826+
&DataType::List(Arc::new(Field::new_list_field(DataType::Int64, false))),
1827+
&[],
1828+
LexOrdering::default(),
1829+
false,
1830+
)?;
1831+
1832+
last_accumulator.update_batch(values)?;
1833+
1834+
Ok(last_accumulator.size())
1835+
}
1836+
1837+
let batch1 = ListArray::from_iter_primitive::<Int32Type, _, _>(
1838+
repeat_with(|| Some(vec![Some(1)])).take(10000),
1839+
);
1840+
let batch2 =
1841+
ListArray::from_iter_primitive::<Int32Type, _, _>([Some(vec![Some(1)])]);
1842+
1843+
let size1 = size_after_batch(&[Arc::new(batch1)])?;
1844+
let size2 = size_after_batch(&[Arc::new(batch2)])?;
1845+
assert_eq!(size1, size2);
1846+
1847+
Ok(())
1848+
}
17751849
}

‎datafusion/functions-aggregate/src/min_max.rs‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -645,19 +645,29 @@ fn min_max_batch_struct(array: &ArrayRef, ordering: Ordering) -> Result<ScalarVa
645645
}
646646
}
647647
}
648-
// use force_clone to free array reference
649-
Ok(extreme.force_clone())
648+
649+
Ok(extreme)
650650
}
651651

652652
macro_rules! min_max_struct {
653653
($VALUE:expr, $DELTA:expr, $OP:ident) => {{
654654
if $VALUE.is_null() {
655-
$DELTA.clone()
655+
let mut delta_copy = $DELTA.clone();
656+
// When the new value won we want to compact it to
657+
// avoid storing the entire input
658+
delta_copy.compact();
659+
delta_copy
656660
} else if $DELTA.is_null() {
657661
$VALUE.clone()
658662
} else {
659663
match $VALUE.partial_cmp(&$DELTA) {
660-
Some(choose_min_max!($OP)) => $DELTA.clone(),
664+
Some(choose_min_max!($OP)) => {
665+
// When the new value won we want to compact it to
666+
// avoid storing the entire input
667+
let mut delta_copy = $DELTA.clone();
668+
delta_copy.compact();
669+
delta_copy
670+
}
661671
_ => $VALUE.clone(),
662672
}
663673
}

0 commit comments

Comments
 (0)