Skip to content

Commit ee11af2

Browse files
robert3005claude
andauthored
perf(compressor): count narrow integer ranges without hashing (#10041)
Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Claude <noreply@anthropic.com>
1 parent 1df09c0 commit ee11af2

1 file changed

Lines changed: 170 additions & 38 deletions

File tree

‎vortex-compressor/src/stats/integer.rs‎

Lines changed: 170 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,17 @@ where
363363
let null_count = validity.false_count();
364364
let value_count = validity.true_count();
365365

366+
let array_ref = array.as_ref();
367+
let min = array_ref
368+
.statistics()
369+
.compute_as::<T>(Stat::Min, ctx)
370+
.vortex_expect("min should be computed");
371+
372+
let max = array_ref
373+
.statistics()
374+
.compute_as::<T>(Stat::Max, ctx)
375+
.vortex_expect("max should be computed");
376+
366377
// Initialize loop state.
367378
let head_idx = validity
368379
.first()
@@ -371,10 +382,10 @@ where
371382
let head = buffer[head_idx];
372383

373384
let mut loop_state = LoopState {
374-
distinct_values: if count_distinct_values {
375-
HashMap::with_capacity_and_hasher(array.len() / 2, FxBuildHasher)
385+
distinct: if count_distinct_values {
386+
DistinctCounter::new(min, max, array.len())
376387
} else {
377-
HashMap::with_hasher(FxBuildHasher)
388+
DistinctCounter::Off
378389
},
379390
prev: head,
380391
pending: 0,
@@ -432,33 +443,7 @@ where
432443
}
433444

434445
let runs = loop_state.runs;
435-
436-
let array_ref = array.as_ref();
437-
let min = array_ref
438-
.statistics()
439-
.compute_as::<T>(Stat::Min, ctx)
440-
.vortex_expect("min should be computed");
441-
442-
let max = array_ref
443-
.statistics()
444-
.compute_as::<T>(Stat::Max, ctx)
445-
.vortex_expect("max should be computed");
446-
447-
let distinct = count_distinct_values.then(|| {
448-
let (&top_value, &top_count) = loop_state
449-
.distinct_values
450-
.iter()
451-
.max_by_key(|&(_, &count)| count)
452-
.vortex_expect("we know this is non-empty");
453-
454-
DistinctInfo {
455-
distinct_count: u32::try_from(loop_state.distinct_values.len())
456-
.vortex_expect("there are more than `u32::MAX` distinct values"),
457-
most_frequent_value: top_value.0,
458-
top_frequency: top_count,
459-
distinct_values: loop_state.distinct_values,
460-
}
461-
});
446+
let distinct = loop_state.distinct.finish();
462447

463448
let typed = TypedStats { min, max, distinct };
464449

@@ -473,30 +458,128 @@ where
473458
})
474459
}
475460

461+
/// Value ranges up to this many values may be counted in a dense array instead of a hash map.
462+
const DENSE_DISTINCT_MAX_RANGE: usize = 1 << 16;
463+
464+
/// Value ranges up to this many values are always counted in a dense array, regardless of the array
465+
/// length. This covers every `u8` and `i8` array.
466+
const DENSE_DISTINCT_ALWAYS_RANGE: usize = 1 << 8;
467+
468+
/// Accumulates the occurrences of each distinct valid value.
469+
enum DistinctCounter<T> {
470+
/// Distinct values are not being counted.
471+
Off,
472+
/// Counts indexed by `value - min`, used when the value range is narrow. This avoids hashing
473+
/// entirely.
474+
Dense {
475+
/// The minimum valid value.
476+
min: T,
477+
/// `min` cast to `usize`. Casting preserves values modulo `2^usize::BITS`, so the wrapping
478+
/// difference of two cast values is the exact difference of any two values in the range.
479+
min_index: usize,
480+
/// The number of occurrences of `min + index`.
481+
counts: Vec<u32>,
482+
},
483+
/// Counts keyed by value.
484+
Hashed(HashMap<NativeValue<T>, u32, FxBuildHasher>),
485+
}
486+
487+
impl<T: IntegerPType> DistinctCounter<T>
488+
where
489+
NativeValue<T>: Eq + Hash,
490+
{
491+
/// Chooses a dense counter when the valid values span a narrow range, otherwise a hash map.
492+
fn new(min: T, max: T, len: usize) -> Self {
493+
let range_len = match (min.to_i128(), max.to_i128()) {
494+
(Some(min), Some(max)) => usize::try_from(max - min + 1).unwrap_or(usize::MAX),
495+
_ => usize::MAX,
496+
};
497+
498+
// A dense counter is only worthwhile when it is not much larger than the array itself.
499+
if range_len <= DENSE_DISTINCT_ALWAYS_RANGE
500+
|| (range_len <= DENSE_DISTINCT_MAX_RANGE && range_len <= len)
501+
{
502+
Self::Dense {
503+
min,
504+
min_index: min.as_(),
505+
counts: vec![0; range_len],
506+
}
507+
} else {
508+
Self::Hashed(HashMap::with_capacity_and_hasher(len / 2, FxBuildHasher))
509+
}
510+
}
511+
512+
/// Records `count` occurrences of a valid value.
513+
#[allow(clippy::inline_always)]
514+
#[inline(always)]
515+
fn add(&mut self, value: T, count: u32) {
516+
match self {
517+
Self::Off => {}
518+
Self::Dense {
519+
min_index, counts, ..
520+
} => counts[value.as_().wrapping_sub(*min_index)] += count,
521+
Self::Hashed(distinct_values) => {
522+
*distinct_values.entry(NativeValue(value)).or_insert(0) += count
523+
}
524+
}
525+
}
526+
527+
/// Returns the distinct value information, or `None` if counting was off.
528+
fn finish(self) -> Option<DistinctInfo<T>> {
529+
let distinct_values: HashMap<NativeValue<T>, u32, FxBuildHasher> = match self {
530+
Self::Off => return None,
531+
Self::Dense { min, counts, .. } => {
532+
let min = min.to_i128().vortex_expect("integers fit in i128");
533+
counts
534+
.iter()
535+
.enumerate()
536+
.filter(|&(_, &count)| count > 0)
537+
.map(|(index, &count)| {
538+
let value = <T as num_traits::NumCast>::from(min + index as i128)
539+
.vortex_expect("values between min and max fit in the type");
540+
(NativeValue(value), count)
541+
})
542+
.collect()
543+
}
544+
Self::Hashed(distinct_values) => distinct_values,
545+
};
546+
547+
let (&top_value, &top_count) = distinct_values
548+
.iter()
549+
.max_by_key(|&(_, &count)| count)
550+
.vortex_expect("we know this is non-empty");
551+
552+
Some(DistinctInfo {
553+
distinct_count: u32::try_from(distinct_values.len())
554+
.vortex_expect("there are more than `u32::MAX` distinct values"),
555+
most_frequent_value: top_value.0,
556+
top_frequency: top_count,
557+
distinct_values,
558+
})
559+
}
560+
}
561+
476562
/// Internal loop state for integer stats computation.
477563
struct LoopState<T> {
478564
/// The previous value seen.
479565
prev: T,
480-
/// Occurrences of `prev` in the current run not yet added to `distinct_values`.
566+
/// Occurrences of `prev` in the current run not yet added to `distinct`.
481567
pending: u32,
482568
/// The run count.
483569
runs: u32,
484-
/// The distinct values map.
485-
distinct_values: HashMap<NativeValue<T>, u32, FxBuildHasher>,
570+
/// The distinct values counter.
571+
distinct: DistinctCounter<T>,
486572
}
487573

488574
impl<T: IntegerPType> LoopState<T>
489575
where
490576
NativeValue<T>: Eq + Hash,
491577
{
492-
/// Adds the pending occurrences of `prev` to the distinct values map.
578+
/// Adds the pending occurrences of `prev` to the distinct values counter.
493579
#[allow(clippy::inline_always)]
494580
#[inline(always)]
495581
fn flush(&mut self) {
496-
*self
497-
.distinct_values
498-
.entry(NativeValue(self.prev))
499-
.or_insert(0) += self.pending;
582+
self.distinct.add(self.prev, self.pending);
500583
self.pending = 0;
501584
}
502585

@@ -686,6 +769,55 @@ mod tests {
686769
Ok(())
687770
}
688771

772+
#[test]
773+
fn dense_distinct_signed_full_range() -> VortexResult<()> {
774+
let mut ctx = array_session().create_execution_ctx();
775+
let values: Buffer<i8> = (i8::MIN..=i8::MAX)
776+
.chain(i8::MIN..=i8::MAX)
777+
.chain(iter::once(i8::MAX))
778+
.collect();
779+
let array = PrimitiveArray::new(values, Validity::NonNullable);
780+
781+
let stats = typed_int_stats::<i8>(&array, true, &mut ctx)?;
782+
assert_eq!(stats.distinct_count(), Some(256));
783+
assert_eq!(
784+
stats.most_frequent_value_and_count(),
785+
Some((i8::MAX.into(), 3))
786+
);
787+
Ok(())
788+
}
789+
790+
#[test]
791+
fn dense_distinct_ignores_values_under_nulls() -> VortexResult<()> {
792+
let mut ctx = array_session().create_execution_ctx();
793+
let array = PrimitiveArray::new(
794+
buffer![1000i32, 1001, 1000, 5_000_000, 1003],
795+
Validity::from_iter([true, true, true, false, true]),
796+
);
797+
798+
let stats = typed_int_stats::<i32>(&array, true, &mut ctx)?;
799+
assert_eq!(stats.distinct_count(), Some(3));
800+
assert_eq!(
801+
stats.most_frequent_value_and_count(),
802+
Some((1000i32.into(), 2))
803+
);
804+
Ok(())
805+
}
806+
807+
#[test]
808+
fn wide_range_distinct() -> VortexResult<()> {
809+
let mut ctx = array_session().create_execution_ctx();
810+
let array = PrimitiveArray::new(buffer![0u64, u64::MAX, 0], Validity::NonNullable);
811+
812+
let stats = typed_int_stats::<u64>(&array, true, &mut ctx)?;
813+
assert_eq!(stats.distinct_count(), Some(2));
814+
assert_eq!(
815+
stats.most_frequent_value_and_count(),
816+
Some((0u64.into(), 2))
817+
);
818+
Ok(())
819+
}
820+
689821
#[test]
690822
fn test_integer_stats_leading_nulls() {
691823
let mut ctx = array_session().create_execution_ctx();

0 commit comments

Comments
 (0)