Skip to content

Commit ce4e7c6

Browse files
rluvatonrluvaton
committed
add unsafe, and always check in bounds
also change the block size to be just batch size instead of the << 18
1 parent 8fc173c commit ce4e7c6

4 files changed

Lines changed: 168 additions & 54 deletions

File tree

‎datafusion/expr-common/src/blocked_groups_accumulator.rs‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,37 @@ impl BlocksIndex {
7878
self.block_index() * block_size + self.index_in_block()
7979
}
8080

81+
/// Returns `true` if every index in `indices` points to an existing group
82+
/// in storage that holds `len` groups in blocks of `block_size`, where
83+
/// every block is full except maybe the last one.
84+
///
85+
/// For example, with `block_size = 4` and `len = 10` the blocks hold 4, 4
86+
/// and 2 groups, so `(2, 1)` is valid, but `(2, 2)`, `(3, 0)` and `(0, 4)`
87+
/// are not.
88+
pub fn all_in_bounds(indices: &[Self], block_size: usize, len: usize) -> bool {
89+
let Some(last_group) = len.checked_sub(1) else {
90+
return indices.is_empty();
91+
};
92+
let last = Self::from_flat(last_group, block_size);
93+
94+
// A single pass without a branch per index, so it is vectorized
95+
let (max, or) = indices
96+
.iter()
97+
.fold((0, 0), |(max, or), index| (max.max(index.0), or | index.0));
98+
99+
// Indices order by block, then position, so an index is valid exactly
100+
// when it is not after the last group and its position is below
101+
// `block_size`
102+
let positions_in_block = if block_size.is_power_of_two() {
103+
(or as u32 as usize) < block_size
104+
} else {
105+
indices
106+
.iter()
107+
.all(|index| index.index_in_block() < block_size)
108+
};
109+
max <= last.0 && positions_in_block
110+
}
111+
81112
/// Writes [`Self::from_flat`] of every index in `flat` to `out`,
82113
/// replacing its contents.
83114
pub fn from_flat_slice(flat: &[usize], block_size: usize, out: &mut Vec<Self>) {
@@ -213,6 +244,34 @@ pub trait BlockedGroupsAccumulator: Send + Any {
213244
mod tests {
214245
use super::*;
215246

247+
#[test]
248+
fn all_in_bounds() {
249+
let in_bounds = |indices: &[(usize, usize)], block_size, len| {
250+
let indices: Vec<_> = indices
251+
.iter()
252+
.map(|&(b, i)| BlocksIndex::new(b, i))
253+
.collect();
254+
BlocksIndex::all_in_bounds(&indices, block_size, len)
255+
};
256+
257+
// Blocks of 4, 4 and 2 groups
258+
assert!(in_bounds(&[(0, 0), (1, 3), (2, 1)], 4, 10));
259+
// After the last group of the last block
260+
assert!(!in_bounds(&[(0, 0), (2, 2)], 4, 10));
261+
// After the last block
262+
assert!(!in_bounds(&[(3, 0)], 4, 10));
263+
// Position past the end of a full block
264+
assert!(!in_bounds(&[(0, 4)], 4, 10));
265+
266+
// Block size that is not a power of two: blocks of 3, 3 and 1 groups
267+
assert!(in_bounds(&[(1, 2), (2, 0)], 3, 7));
268+
assert!(!in_bounds(&[(0, 3)], 3, 7));
269+
270+
// No groups
271+
assert!(in_bounds(&[], 4, 0));
272+
assert!(!in_bounds(&[(0, 0)], 4, 0));
273+
}
274+
216275
#[test]
217276
fn from_emit_to_splits_into_blocks() {
218277
assert_eq!(

‎datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/blocked_vec.rs‎

Lines changed: 81 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -232,27 +232,53 @@ impl<T: Copy, const FIXED_BLOCK_SIZE: bool> BlockedVec<T, FIXED_BLOCK_SIZE> {
232232
/// `group_indices` must only point to the first `total_num_groups`
233233
/// groups, which is the contract of
234234
/// [`BlockedGroupsAccumulator::update_batch`] and
235-
/// [`BlockedGroupsAccumulator::merge_batch`]. It is only checked in debug
236-
/// builds, like the flat accumulators' unchecked indexing.
235+
/// [`BlockedGroupsAccumulator::merge_batch`].
237236
///
238237
/// [`BlockedGroupsAccumulator::update_batch`]: datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator::update_batch
239238
/// [`BlockedGroupsAccumulator::merge_batch`]: datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator::merge_batch
239+
///
240+
/// # Panics
241+
/// If an index in `group_indices` is not below `total_num_groups`.
242+
/// if there is an index in block that is larger than block size
240243
pub fn update<F>(
241244
&mut self,
242245
total_num_groups: usize,
243246
starting_value: T,
244247
group_indices: &[BlocksIndex],
245248
nulls: Option<&NullBuffer>,
246249
opt_filter: Option<&BooleanArray>,
247-
mut update_fn: F,
250+
update_fn: F,
248251
) where
249252
F: FnMut(&mut T),
250253
{
251254
self.grow_to(total_num_groups, starting_value);
252-
self.debug_assert_in_bounds(group_indices);
255+
self.assert_in_bounds(group_indices);
256+
257+
// SAFETY: grown to `total_num_groups` and checked above
258+
unsafe {
259+
self.update_unchecked(group_indices, nulls, opt_filter, update_fn);
260+
}
261+
}
262+
263+
/// Same as [`Self::update`], but does not grow the vector and the caller guarantees that indices are in bounds.
264+
///
265+
/// # Safety
266+
/// The behavior is undefined if any of the following are true:
267+
/// 1. `groups_indices` contain block index outside the number of blocks (e.g. you did not call `grow_to` first to the expected number of values)
268+
/// 2. `groups_indices` index in block is larger than block size
269+
/// 3. `groups_indices` index in block for last blocks is larger than the number of elements in the last block
270+
pub unsafe fn update_unchecked<F>(
271+
&mut self,
272+
group_indices: &[BlocksIndex],
273+
nulls: Option<&NullBuffer>,
274+
opt_filter: Option<&BooleanArray>,
275+
mut update_fn: F,
276+
) where
277+
F: FnMut(&mut T),
278+
{
253279
if let Some(block) = self.as_single_block_mut() {
254280
accumulate_blocked_indices(group_indices, nulls, opt_filter, |index| {
255-
// SAFETY: indices are in bounds, see the method docs
281+
// SAFETY: indices are in bounds, guaranteed by the caller
256282
update_fn(unsafe { block.get_unchecked_mut(index.index_in_block()) })
257283
});
258284
return;
@@ -263,7 +289,7 @@ impl<T: Copy, const FIXED_BLOCK_SIZE: bool> BlockedVec<T, FIXED_BLOCK_SIZE> {
263289
let mut addrs = [std::ptr::null_mut::<T>(); RESOLVE_CHUNK];
264290
for chunk in group_indices.chunks(RESOLVE_CHUNK) {
265291
for (addr, &index) in addrs.iter_mut().zip(chunk) {
266-
// SAFETY: indices are in bounds, see the method docs
292+
// SAFETY: indices are in bounds, guaranteed by the caller
267293
*addr = unsafe { ptrs.ptr(index) };
268294
}
269295
for &addr in &addrs[..chunk.len()] {
@@ -274,7 +300,7 @@ impl<T: Copy, const FIXED_BLOCK_SIZE: bool> BlockedVec<T, FIXED_BLOCK_SIZE> {
274300
}
275301
} else {
276302
accumulate_blocked_indices(group_indices, nulls, opt_filter, |index| {
277-
// SAFETY: indices are in bounds, see the method docs
303+
// SAFETY: indices are in bounds, guaranteed by the caller
278304
update_fn(unsafe { &mut *ptrs.ptr(index) })
279305
});
280306
}
@@ -289,31 +315,52 @@ impl<T: Copy, const FIXED_BLOCK_SIZE: bool> BlockedVec<T, FIXED_BLOCK_SIZE> {
289315
/// `group_indices` must only point to the first `total_num_groups`
290316
/// groups, which is the contract of
291317
/// [`BlockedGroupsAccumulator::update_batch`] and
292-
/// [`BlockedGroupsAccumulator::merge_batch`]. It is only checked in debug
293-
/// builds, like the flat accumulators' unchecked indexing.
318+
/// [`BlockedGroupsAccumulator::merge_batch`].
294319
///
295320
/// [`BlockedGroupsAccumulator::update_batch`]: datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator::update_batch
296321
/// [`BlockedGroupsAccumulator::merge_batch`]: datafusion_expr_common::blocked_groups_accumulator::BlockedGroupsAccumulator::merge_batch
297322
///
298323
/// # Panics
299-
/// If `values` and `group_indices` have different lengths.
324+
/// If `values` and `group_indices` have different lengths, or an index in
325+
/// `group_indices` is not below `total_num_groups`.
300326
pub fn update_with<V, F>(
301327
&mut self,
302328
total_num_groups: usize,
303329
starting_value: T,
304330
group_indices: &[BlocksIndex],
305331
values: &[V],
332+
update_fn: F,
333+
) where
334+
V: Copy,
335+
F: FnMut(&mut T, V),
336+
{
337+
self.grow_to(total_num_groups, starting_value);
338+
self.assert_in_bounds(group_indices);
339+
340+
// SAFETY: grown to `total_num_groups` and checked above
341+
unsafe { self.update_with_unchecked(group_indices, values, update_fn) }
342+
}
343+
344+
/// Same as [`Self::update_with`], but does not grow the vector and the caller guarantees that indices are in bounds.
345+
///
346+
/// # Safety
347+
/// The behavior is undefined if any of the following are true:
348+
/// 1. `groups_indices` contain block index outside the number of blocks (e.g. you did not call `grow_to` first to the expected number of values)
349+
/// 2. `groups_indices` index in block is larger than block size
350+
/// 3. `groups_indices` index in block for last block is larger than the number of elements in the last block (e.g. you did not call `grow_to` first to the expected number of values)
351+
pub unsafe fn update_with_unchecked<V, F>(
352+
&mut self,
353+
group_indices: &[BlocksIndex],
354+
values: &[V],
306355
mut update_fn: F,
307356
) where
308357
V: Copy,
309358
F: FnMut(&mut T, V),
310359
{
311360
assert_eq!(group_indices.len(), values.len());
312-
self.grow_to(total_num_groups, starting_value);
313-
self.debug_assert_in_bounds(group_indices);
314361
if let Some(block) = self.as_single_block_mut() {
315362
for (&index, &value) in group_indices.iter().zip(values) {
316-
// SAFETY: indices are in bounds, see the method docs
363+
// SAFETY: indices are in bounds, guaranteed by the caller
317364
update_fn(
318365
unsafe { block.get_unchecked_mut(index.index_in_block()) },
319366
value,
@@ -329,7 +376,7 @@ impl<T: Copy, const FIXED_BLOCK_SIZE: bool> BlockedVec<T, FIXED_BLOCK_SIZE> {
329376
.zip(values.chunks(RESOLVE_CHUNK))
330377
{
331378
for (addr, &index) in addrs.iter_mut().zip(indices) {
332-
// SAFETY: indices are in bounds, see the method docs
379+
// SAFETY: indices are in bounds, guaranteed by the caller
333380
*addr = unsafe { ptrs.ptr(index) };
334381
}
335382
for (&addr, &value) in addrs.iter().zip(values) {
@@ -340,15 +387,13 @@ impl<T: Copy, const FIXED_BLOCK_SIZE: bool> BlockedVec<T, FIXED_BLOCK_SIZE> {
340387
}
341388
}
342389

343-
/// Checks, in debug builds only, that every index points to an element.
390+
/// Checks that every index points to an element, so the update loops
391+
/// can index without bounds checks. Every block but the last holds
392+
/// `block_size` elements, see [`BlocksIndex::all_in_bounds`].
344393
#[inline]
345-
fn debug_assert_in_bounds(&self, group_indices: &[BlocksIndex]) {
346-
debug_assert!(
347-
group_indices.iter().all(|index| {
348-
self.blocks
349-
.get(index.block_index())
350-
.is_some_and(|block| index.index_in_block() < block.len())
351-
}),
394+
fn assert_in_bounds(&self, group_indices: &[BlocksIndex]) {
395+
assert!(
396+
BlocksIndex::all_in_bounds(group_indices, self.block_size, self.len),
352397
"group index out of bounds"
353398
);
354399
}
@@ -673,4 +718,18 @@ mod tests {
673718
fn zero_block_size_panics() {
674719
let _ = BlockedVec::<i64>::new(0);
675720
}
721+
722+
#[test]
723+
#[should_panic(expected = "group index out of bounds")]
724+
fn update_out_of_bounds_panics() {
725+
let mut v = BlockedVec::<u64>::new(4);
726+
v.update(1, 0, &[BlocksIndex::new(0, 1)], None, None, |v| *v += 1);
727+
}
728+
729+
#[test]
730+
#[should_panic(expected = "group index out of bounds")]
731+
fn update_with_out_of_bounds_panics() {
732+
let mut v = BlockedVec::<u64>::new(4);
733+
v.update_with(6, 0, &[BlocksIndex::new(0, 4)], &[1], |v, x| *v += x);
734+
}
676735
}

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

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -892,16 +892,18 @@ impl BlockedGroupsAccumulator for BlockedCountGroupsAccumulator {
892892
let values = &values[0];
893893
let nulls = values.logical_nulls().filter(|n| n.null_count() > 0);
894894

895-
// Add one to each group's counter for each non null, non
896-
// filtered value
897-
self.counts.update(
898-
total_num_groups,
899-
0,
900-
group_indices,
901-
nulls.as_ref(),
902-
opt_filter,
903-
|count| *count += 1,
904-
);
895+
self.counts.grow_to(total_num_groups, 0);
896+
897+
// Add one to each group's counter for each non null, non filtered value
898+
// SAFETY: group_index is guaranteed to be in bounds and less than total_num_groups
899+
unsafe {
900+
self.counts.update_unchecked(
901+
group_indices,
902+
nulls.as_ref(),
903+
opt_filter,
904+
|count| *count += 1,
905+
);
906+
}
905907
Ok(())
906908
}
907909

‎datafusion/physical-plan/src/aggregates/aggregate_hash_table/storage.rs‎

Lines changed: 16 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -60,15 +60,11 @@ use crate::aggregates::order::GroupOrdering;
6060

6161
use super::common::{AggregateAccumulator, MaterializeAccumulatorFn};
6262

63-
/// Smallest block, in groups. Tables with fewer groups stay in a single,
64-
/// flat block and run exactly the flat hot loops.
65-
const MIN_BLOCK_SIZE: usize = 1 << 18;
66-
6763
/// Number of groups per block for every blocked part of a table.
6864
///
69-
/// At least `batch_size`, so an output batch never spans two blocks.
65+
/// `batch_size`, so every block is emitted as one output batch.
7066
pub(in crate::aggregates) fn block_size(batch_size: usize) -> usize {
71-
batch_size.max(MIN_BLOCK_SIZE).next_power_of_two()
67+
batch_size
7268
}
7369

7470
/// Returns `true` if the group keys of `agg` can be stored in blocks: a single
@@ -476,23 +472,22 @@ mod tests {
476472
assert!(!state.accumulators[1].storage().is_blocked());
477473
// min needs flat indices while the keys produce blocked ones
478474
assert!(state.keys.convert);
479-
assert_eq!(block_size(8192), MIN_BLOCK_SIZE);
480-
// The block always holds at least one output batch.
481-
assert_eq!(block_size(MIN_BLOCK_SIZE + 1), MIN_BLOCK_SIZE * 2);
475+
assert_eq!(block_size(8192), 8192);
482476
Ok(())
483477
}
484478

485479
/// Blocked `count` and flat `min` over blocked keys: flat columns are
486480
/// sliced at the block boundaries, and their memory is counted once.
487481
#[test]
488482
fn mixed_storage_output_is_aligned_and_counted_once() -> Result<()> {
489-
let num_groups = MIN_BLOCK_SIZE + 10;
483+
let batch_size = 100;
484+
let num_groups = batch_size + 10;
490485
let input = RecordBatch::try_new(
491486
schema(),
492487
vec![Arc::new(Int64Array::from_iter_values(0..num_groups as i64))],
493488
)?;
494489
let agg = aggregate_exec(vec![], &[count_udaf(), min_udaf()])?;
495-
let mut table = single_table(&agg, 8192)?;
490+
let mut table = single_table(&agg, batch_size)?;
496491
table.aggregate_batch(&input)?;
497492

498493
let batches = table.take_state_batches()?;
@@ -501,7 +496,7 @@ mod tests {
501496
.iter()
502497
.map(|b| b.batch.num_rows())
503498
.collect::<Vec<_>>(),
504-
vec![MIN_BLOCK_SIZE, 10]
499+
vec![batch_size, 10]
505500
);
506501
for b in &batches {
507502
let keys = b.batch.column(0).as_primitive::<Int64Type>().values();
@@ -527,8 +522,8 @@ mod tests {
527522
/// released once all its slices are handed out.
528523
#[test]
529524
fn output_releases_state_block_by_block() -> Result<()> {
530-
let num_groups = MIN_BLOCK_SIZE + 10;
531-
let batch_size = 8192;
525+
let batch_size = 100;
526+
let num_groups = 4 * batch_size + 10;
532527
let input = RecordBatch::try_new(
533528
schema(),
534529
vec![Arc::new(Int64Array::from_iter_values(0..num_groups as i64))],
@@ -563,15 +558,14 @@ mod tests {
563558

564559
keys.sort_unstable();
565560
assert_eq!(keys, (0..num_groups as i64).collect::<Vec<_>>());
566-
// While the first block is sliced, the table still holds the second
567-
// block; once the first block is handed out, only the small second
568-
// block (10 groups) remains.
569-
let first_block_batches = MIN_BLOCK_SIZE / batch_size;
570-
assert!(memory_per_batch[0] <= memory_while_building);
561+
// Each output batch is one block, released once handed out, so only
562+
// the last, small block (10 groups) remains after the full blocks.
563+
assert_eq!(memory_per_batch.len(), 5);
564+
assert!(memory_per_batch.is_sorted_by(|a, b| a > b));
571565
assert!(
572-
memory_per_batch[first_block_batches] < memory_while_building / 4,
573-
"memory after the first block: {}, while building: {memory_while_building}",
574-
memory_per_batch[first_block_batches]
566+
memory_per_batch[3] < memory_while_building / 4,
567+
"memory after the full blocks: {}, while building: {memory_while_building}",
568+
memory_per_batch[3]
575569
);
576570
assert_eq!(*memory_per_batch.last().unwrap(), 0);
577571
Ok(())

0 commit comments

Comments
 (0)