Skip to content

perf: reuse cached dictionary value slots in vectorized_equal_to #25652

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jayzhan211 merged 3 commits into from
Sep 30, 2026
Merged

perf: reuse cached dictionary value slots in vectorized_equal_to #25652

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f75c111
perf: reuse cached dictionary value slots in vectorized_equal_to
CuteChuanChuan Sep 23, 2026
f98cc4a
Merge branch 'main' into issue-25219/reuse-cache-dic
CuteChuanChuan Sep 26, 2026
40bfa2b
perf: address review feedback on vectorized_equal_to
CuteChuanChuan Sep 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
236 changes: 182 additions & 54 deletions datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -232,6 +232,42 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> DictionaryGroupValuesColumn<K> {
}
}

/// Compares the value at `val_idx` against `inner[lhs_slot]` directly,
/// memorizing `lhs_slot` in `val_to_inner` when they are equal.
///
/// Each distinct value occupies exactly one slot in `inner`, so an equal
/// comparison proves `lhs_slot` is the slot for `val_idx` without probing
/// `value_dedup` or hashing. An unequal result is deliberately not
/// memoized: it says which slot the value is *not* in, not which one it
/// is. That is what keeps `val_to_inner` fill-only and never stale.
///
/// Callers must have run [`Self::sync_value_cache`] for `dict_values`.
///
/// Only called on a `val_to_inner` miss. It takes the fields it needs
/// rather than `&mut self`, so the caller's hot loop can keep its slices
/// of `val_to_inner` and `group_to_inner` in registers across the call.
#[cold]
#[inline(never)]
fn equal_to_uncached(
val_to_inner: &mut [usize],
inner: &dyn GroupColumn,
null_array: &ArrayRef,
dict_values: &ArrayRef,
val_idx: usize,
lhs_slot: usize,
) -> bool {
let equal = if dict_values.is_null(val_idx) {
inner.equal_to(lhs_slot, null_array, 0)
} else {
inner.equal_to(lhs_slot, dict_values, val_idx)
};

if equal {
val_to_inner[val_idx] = lhs_slot;
}
equal
}

fn find_or_insert_null(&mut self) -> Result<usize> {
if let Some(slot) = self.null_inner_slot {
return Ok(slot);
Expand All @@ -242,38 +278,13 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> DictionaryGroupValuesColumn<K> {
Ok(slot)
}

fn build_lookup_table(
&self,
dict_values: &ArrayRef,
val_hashes: &[u64],
) -> Vec<usize> {
let num_distinct = dict_values.len();
let mut table = vec![usize::MAX; num_distinct + 1];
let inner = &*self.inner;
for val_idx in 0..num_distinct {
if dict_values.is_null(val_idx) {
table[val_idx] = self.null_inner_slot.unwrap_or(usize::MAX);
} else {
let hash = val_hashes[val_idx];
if let Some(&(_, slot)) =
self.value_dedup.find(hash, |&(entry_hash, slot)| {
entry_hash == hash && inner.equal_to(slot, dict_values, val_idx)
})
{
table[val_idx] = slot;
}
}
}
table[num_distinct] = self.null_inner_slot.unwrap_or(usize::MAX);
table
}

/// Per-row fallback for `vectorized_equal_to` used when the number of rows
/// to check is smaller than the dictionary cardinality, making the O(D)
/// lookup-table build more expensive than direct value comparison.
/// Per-row fallback for `vectorized_equal_to` used when the values array is
/// not cached and the number of rows to check is smaller than the
/// dictionary cardinality, making the O(D) hashing in `sync_value_cache`
/// more expensive than direct value comparison.
///
/// `#[cold]` + `#[inline(never)]` keeps this code out of the hot
/// lookup-table loops in `vectorized_equal_to` so LLVM can pipeline them.
/// `val_to_inner` loops in `vectorized_equal_to` so LLVM can pipeline them.
#[cold]
#[inline(never)]
fn equal_to_per_row(
Expand Down Expand Up @@ -348,12 +359,12 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
let dict = array.as_dictionary::<K>();
let dict_keys = dict.keys();
let dict_values = dict.values();
let num_distinct = dict_values.len();

// The fallback is in a separate #[cold] function so its code does not
// appear inline here and cannot prevent LLVM from pipelining / unrolling
// the hot lookup-table loops below.
if rhs_rows.len() < num_distinct {
let cached = self
.cached_values
.as_ref()
.is_some_and(|c| Arc::ptr_eq(c, dict_values));
if !cached && rhs_rows.len() < dict_values.len() {
self.equal_to_per_row(
lhs_rows,
dict_values,
Expand All @@ -363,44 +374,64 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
);
return;
}
self.sync_value_cache(dict_values);

@jayzhan211 jayzhan211 Sep 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dropping the rhs_rows.len() < num_distinct fallback regresses the steady-state case with a new values array per batch: every row hits an existing group, so vectorized_append isn't called for the column. sync_value_cache now hashes all D values per batch where base did per-row equal_to. That happens with concatenated/coalesced dictionaries and at parquet row-group boundaries. With 8192 rows/batch and D = 100000, I measured 1.19–1.48x vs base. Keep the fallback, but only when the Arc isn't already cached (0.89–0.95x in my run):

let cached = self
    .cached_values
    .as_ref()
    .is_some_and(|c| Arc::ptr_eq(c, dict_values));
if !cached && rhs_rows.len() < dict_values.len() {
    self.equal_to_per_row(lhs_rows, dict_values, dict, rhs_rows, equal_to_results);
    return;
}
self.sync_value_cache(dict_values);

Please also add before/after bench numbers to the PR description, since this is a perf PR.

@CuteChuanChuan CuteChuanChuan Sep 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done, using your condition. The fallback now only runs when the values Arc isn't cached and rhs_rows.len() < dict_values.len().


let mut val_hashes = vec![0u64; dict_values.len()];
create_hashes(
std::slice::from_ref(dict_values),
&self.random_state,
&mut val_hashes,
)
.unwrap();
let lookup = self.build_lookup_table(dict_values, &val_hashes);

let raw_keys = dict_keys.values();
let group_to_inner = self.group_to_inner.as_slice();

let val_to_inner = self.val_to_inner.as_mut_slice();
let inner = &*self.inner;
let null_array = &self.null_array;
let null_inner_slot = self.null_inner_slot;
if dict_keys.null_count() == 0 {
// No null keys : skip the get_bit guard: we only ever write false,
// so overwriting an already-false bit is a no-op.
let raw_keys = dict_keys.values();
for (idx, (&lhs_row, &rhs_row)) in
lhs_rows.iter().zip(rhs_rows.iter()).enumerate()

@jayzhan211 jayzhan211 Sep 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Calling lookup_inner_slot(&mut self, ..) on every row makes the hot loop slower than the old lookup[..] slice index. The existing dictionary_group_values bench regresses: dict_intern_emit 1.21–1.71x at card ≤ 1000, dict_repeated_intern_emit 1.08–1.52x. In intern_emit the values Arc is already cached by vectorized_append, so this is the loop itself, not the removed fallback. Reading the cache inline and taking the resolver only on a miss brought intern_emit back to 0.96–1.05x in my run:

-                let val_idx = raw_keys[rhs_row].as_usize();
-                let rhs_slot = self.lookup_inner_slot(dict_values, val_idx);
+                let val_idx = raw_keys[rhs_row].as_usize();
+                let rhs_slot = match self.val_to_inner[val_idx] {
+                    usize::MAX => self.lookup_inner_slot(dict_values, val_idx),
+                    slot => slot,
+                };
+    #[cold]
+    #[inline(never)]
     fn lookup_inner_slot(&mut self, dict_values: &ArrayRef, val_idx: usize) -> usize {

Apply the same change to the nullable-key loop and keep its get_bit guard. A merged single loop still left repeated_intern_emit (10% null keys) at 1.06–1.19x.

@CuteChuanChuan CuteChuanChuan Sep 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the numbers. I applied this and then ran the whole dictionary_group_values bench against main.

{
let rhs_slot = lookup[raw_keys[rhs_row].as_usize()];
if rhs_slot == usize::MAX || group_to_inner[lhs_row] != rhs_slot {
let val_idx = raw_keys[rhs_row].as_usize();
let lhs_slot = group_to_inner[lhs_row];
let equal = match val_to_inner[val_idx] {
usize::MAX => Self::equal_to_uncached(
val_to_inner,
inner,
null_array,
dict_values,
val_idx,
lhs_slot,
),
slot => slot == lhs_slot,
};
if !equal {
equal_to_results.set_bit(idx, false);
}
}
} else {
let null_buf = dict_keys.nulls().unwrap();
let raw_keys = dict_keys.values();
for (idx, (&lhs_row, &rhs_row)) in
lhs_rows.iter().zip(rhs_rows.iter()).enumerate()
{
if equal_to_results.get_bit(idx) {
let val_idx = if null_buf.is_null(rhs_row) {
num_distinct
// A null key is not a position in the values array, so it
// resolves straight to the null slot instead of going
// through `val_to_inner`.
let lhs_slot = group_to_inner[lhs_row];
let equal = if null_buf.is_null(rhs_row) {
null_inner_slot == Some(lhs_slot)
} else {
raw_keys[rhs_row].as_usize()
let val_idx = raw_keys[rhs_row].as_usize();
match val_to_inner[val_idx] {
usize::MAX => Self::equal_to_uncached(
val_to_inner,
inner,
null_array,
dict_values,
val_idx,
lhs_slot,
),
slot => slot == lhs_slot,
}
};
let rhs_slot = lookup[val_idx];
if rhs_slot == usize::MAX || group_to_inner[lhs_row] != rhs_slot {
if !equal {
equal_to_results.set_bit(idx, false);
}
}
Expand Down Expand Up @@ -1060,4 +1091,101 @@ mod tests {
assert_eq!(col.inner.len(), 2, "'a' must not be duplicated");
assert_eq!(col.len(), 3);
}

/// On a cache miss `vectorized_equal_to` compares against the group's own
/// slot: an equal result proves the slot and is cached, an unequal one
/// only rules that slot out and must leave the entry unresolved.
#[test]
fn vectorized_equal_to_caches_only_equal_results() {
let mut col = utf8_col();
col.vectorized_append(
&i32_dict(&[Some(0), Some(1)], &[Some("a"), Some("b")]),
&[0, 1],
)
.unwrap();
let (group_a, group_b) = (0, 1);
let slot_a = col.group_to_inner[group_a];
let slot_b = col.group_to_inner[group_b];

// A new values array in the opposite order, so `val_idx` != slot.
// Two rows for two values keeps the per-row fallback out of the way.
let (idx_b, idx_a) = (0, 1);
let values: ArrayRef = Arc::new(StringArray::from(vec![Some("b"), Some("a")]));
let a_then_b =
dict_with_values(&[Some(idx_a as i32), Some(idx_b as i32)], &values);

// "a" vs group a is equal, "b" vs group a is not.
let mut buf = all_true(2);
col.vectorized_equal_to(&[group_a, group_a], &a_then_b, &[0, 1], &mut buf);
assert_eq!(bool_vec(&buf), vec![true, false]);
assert_eq!(col.val_to_inner[idx_a], slot_a);
assert_eq!(
col.val_to_inner[idx_b],
usize::MAX,
"unequal result was cached"
);

// Later, "b" vs group b fills the entry that was left unresolved.
let only_b = dict_with_values(&[Some(idx_b as i32)], &values);
let mut buf = all_true(1);
col.vectorized_equal_to(&[group_b], &only_b, &[0], &mut buf);
assert_eq!(bool_vec(&buf), vec![true]);
assert_eq!(col.val_to_inner[idx_b], slot_b);
}

/// A non-null key pointing at a null value and a null key must both
/// compare equal to the null group, and only the former goes through
/// `val_to_inner`.
#[test]
fn vectorized_equal_to_matches_null_value_and_null_key_to_null_group() {
let mut col = utf8_col();
col.vectorized_append(&i32_dict(&[None, Some(0)], &[Some("a")]), &[0, 1])
.unwrap();
let (null_group, group_a) = (0, 1);
let null_slot = col.null_inner_slot.unwrap();

let idx_null = 0;
let values: ArrayRef = Arc::new(StringArray::from(vec![None, Some("a")]));
let batch = dict_with_values(
&[Some(idx_null as i32), Some(idx_null as i32), None],
&values,
);

// null value vs null group, null value vs group a, null key vs null group.
let mut buf = all_true(3);
col.vectorized_equal_to(
&[null_group, group_a, null_group],
&batch,
&[0, 1, 2],
&mut buf,
);
assert_eq!(bool_vec(&buf), vec![true, false, true]);
assert_eq!(col.val_to_inner[idx_null], null_slot);
}

/// Fewer rows than values in an uncached values array takes the per-row
/// fallback, which must not rebuild the cache for that array.
#[test]
fn vectorized_equal_to_small_batch_of_new_values_skips_cache_rebuild() {
let old_values: ArrayRef = Arc::new(StringArray::from(vec![Some("a")]));
let mut col = utf8_col();
col.vectorized_append(&dict_with_values(&[Some(0)], &old_values), &[0])
.unwrap();
let group_a = 0;

// One row against three values in an array the cache has not seen.
let new_values: ArrayRef =
Arc::new(StringArray::from(vec![Some("x"), Some("a"), Some("y")]));
let idx_a = 1;
let only_a = dict_with_values(&[Some(idx_a)], &new_values);
let mut buf = all_true(1);
col.vectorized_equal_to(&[group_a], &only_a, &[0], &mut buf);
assert_eq!(bool_vec(&buf), vec![true]);
assert!(
col.cached_values
.as_ref()
.is_some_and(|c| Arc::ptr_eq(c, &old_values)),
"fallback must leave the cache on the previous values array"
);
}
}
Loading