@@ -517,12 +517,122 @@ fn hash_generic_byte_view_array<T: ByteViewType>(
517517 }
518518}
519519
520+ /// Dense-path unroll width; 8 lanes keep the gather pipeline full on ARM64/x86.
521+ #[ cfg( not( feature = "force_hash_collisions" ) ) ]
522+ const DICT_SCATTER_LANES : usize = 8 ;
523+
524+ #[ cfg( not( feature = "force_hash_collisions" ) ) ]
525+ #[ inline( always) ]
526+ fn apply_dict_hash < const MULTI_COL : bool > ( prev : u64 , dict_hash : u64 ) -> u64 {
527+ if MULTI_COL {
528+ combine_hashes ( dict_hash, prev)
529+ } else {
530+ dict_hash
531+ }
532+ }
533+
534+ /// Scatter dict hashes into every key position, 8-way unrolled.
535+ #[ cfg( not( feature = "force_hash_collisions" ) ) ]
536+ #[ inline( always) ]
537+ fn scatter_dict_dense < K : ArrowDictionaryKeyType , const MULTI_COL : bool > (
538+ keys : & [ K :: Native ] ,
539+ dict_hashes : & [ u64 ] ,
540+ hashes : & mut [ u64 ] ,
541+ ) {
542+ let full = ( keys. len ( ) / DICT_SCATTER_LANES ) * DICT_SCATTER_LANES ;
543+ let ( keys_full, keys_tail) = keys. split_at ( full) ;
544+ let ( hashes_full, hashes_tail) = hashes. split_at_mut ( full) ;
545+ for ( k_chunk, h_chunk) in keys_full
546+ . chunks_exact ( DICT_SCATTER_LANES )
547+ . zip ( hashes_full. chunks_exact_mut ( DICT_SCATTER_LANES ) )
548+ {
549+ let dh: [ u64 ; DICT_SCATTER_LANES ] = std:: array:: from_fn ( |i| unsafe {
550+ * dict_hashes. get_unchecked ( k_chunk[ i] . as_usize ( ) )
551+ } ) ;
552+ for i in 0 ..DICT_SCATTER_LANES {
553+ h_chunk[ i] = apply_dict_hash :: < MULTI_COL > ( h_chunk[ i] , dh[ i] ) ;
554+ }
555+ }
556+ for ( k, h) in keys_tail. iter ( ) . zip ( hashes_tail. iter_mut ( ) ) {
557+ let dh = unsafe { * dict_hashes. get_unchecked ( k. as_usize ( ) ) } ;
558+ * h = apply_dict_hash :: < MULTI_COL > ( * h, dh) ;
559+ }
560+ }
561+
562+ /// Scatter dict hashes but skip positions whose dict value is null.
563+ #[ cfg( not( feature = "force_hash_collisions" ) ) ]
564+ #[ inline( always) ]
565+ fn scatter_dict_dense_filter_values < K : ArrowDictionaryKeyType , const MULTI_COL : bool > (
566+ keys : & [ K :: Native ] ,
567+ dict_values : & dyn Array ,
568+ dict_hashes : & [ u64 ] ,
569+ hashes : & mut [ u64 ] ,
570+ ) {
571+ for ( k, h) in keys. iter ( ) . zip ( hashes. iter_mut ( ) ) {
572+ let idx = k. as_usize ( ) ;
573+ if dict_values. is_valid ( idx) {
574+ let dh = unsafe { * dict_hashes. get_unchecked ( idx) } ;
575+ * h = apply_dict_hash :: < MULTI_COL > ( * h, dh) ;
576+ }
577+ }
578+ }
579+
580+ /// Blend scatter inside a mixed 64-key chunk governed by `mask`.
581+ ///
582+ /// Branchless fast path when dict values are all valid: clamp garbage key
583+ /// indices (from null-key positions), gather dict hashes for every lane, then
584+ /// csel-blend using the mask so only valid positions are written.
585+ #[ cfg( not( feature = "force_hash_collisions" ) ) ]
586+ #[ inline( always) ]
587+ fn scatter_dict_blend_chunk <
588+ K : ArrowDictionaryKeyType ,
589+ const HAS_NULL_VALUES : bool ,
590+ const MULTI_COL : bool ,
591+ > (
592+ keys : & [ K :: Native ] ,
593+ dict_values : & dyn Array ,
594+ dict_hashes : & [ u64 ] ,
595+ hashes : & mut [ u64 ] ,
596+ mask : u64 ,
597+ ) {
598+ debug_assert_eq ! ( keys. len( ) , 64 ) ;
599+ debug_assert_eq ! ( hashes. len( ) , 64 ) ;
600+ if HAS_NULL_VALUES {
601+ let mut remaining = mask;
602+ while remaining != 0 {
603+ let i = remaining. trailing_zeros ( ) as usize ;
604+ remaining &= remaining - 1 ;
605+ let idx = unsafe { keys. get_unchecked ( i) . as_usize ( ) } ;
606+ if dict_values. is_valid ( idx) {
607+ let dh = unsafe { * dict_hashes. get_unchecked ( idx) } ;
608+ let slot = unsafe { hashes. get_unchecked_mut ( i) } ;
609+ * slot = apply_dict_hash :: < MULTI_COL > ( * slot, dh) ;
610+ }
611+ }
612+ return ;
613+ }
614+ let clamp = dict_hashes. len ( ) . saturating_sub ( 1 ) ;
615+ for lane_base in ( 0 ..64 ) . step_by ( DICT_SCATTER_LANES ) {
616+ let dh: [ u64 ; DICT_SCATTER_LANES ] = std:: array:: from_fn ( |i| {
617+ let raw = unsafe { keys. get_unchecked ( lane_base + i) . as_usize ( ) } ;
618+ let idx = raw. min ( clamp) ;
619+ unsafe { * dict_hashes. get_unchecked ( idx) }
620+ } ) ;
621+ for i in 0 ..DICT_SCATTER_LANES {
622+ let slot = unsafe { hashes. get_unchecked_mut ( lane_base + i) } ;
623+ let new = apply_dict_hash :: < MULTI_COL > ( * slot, dh[ i] ) ;
624+ let valid = ( ( mask >> ( lane_base + i) ) & 1 ) == 1 ;
625+ * slot = if valid { new } else { * slot } ;
626+ }
627+ }
628+ }
629+
520630/// Scatter precomputed dictionary value hashes to key positions.
521631///
522- /// Uses const generics to eliminate runtime branching in the hot loop :
523- /// - `HAS_NULL_KEYS`: Whether to check for null dictionary keys
524- /// - `HAS_NULL_VALUES`: Whether to check for null dictionary values
525- /// - `MULTI_COL`: Whether to combine with existing hash (true) or initialize (false)
632+ /// Const generics eliminate runtime branching:
633+ /// - `HAS_NULL_KEYS`: dictionary has null keys to skip
634+ /// - `HAS_NULL_VALUES`: dictionary values may be null (per-index check)
635+ /// - `MULTI_COL`: combine with existing hash vs. overwrite
526636#[ cfg( not( feature = "force_hash_collisions" ) ) ]
527637#[ inline( never) ]
528638fn hash_dictionary_scatter <
@@ -536,28 +646,68 @@ fn hash_dictionary_scatter<
536646 hashes_buffer : & mut [ u64 ] ,
537647) {
538648 let dict_values = array. values ( ) ;
539- if HAS_NULL_KEYS {
540- for ( hash, key) in hashes_buffer. iter_mut ( ) . zip ( array. keys ( ) . iter ( ) ) {
541- if let Some ( key) = key {
542- let idx = key. as_usize ( ) ;
543- if !HAS_NULL_VALUES || dict_values. is_valid ( idx) {
544- if MULTI_COL {
545- * hash = combine_hashes ( dict_hashes[ idx] , * hash) ;
546- } else {
547- * hash = dict_hashes[ idx] ;
548- }
549- }
649+ let keys = array. keys ( ) . values ( ) ;
650+
651+ if !HAS_NULL_KEYS {
652+ if HAS_NULL_VALUES {
653+ scatter_dict_dense_filter_values :: < K , MULTI_COL > (
654+ keys,
655+ dict_values. as_ref ( ) ,
656+ dict_hashes,
657+ hashes_buffer,
658+ ) ;
659+ } else {
660+ scatter_dict_dense :: < K , MULTI_COL > ( keys, dict_hashes, hashes_buffer) ;
661+ }
662+ return ;
663+ }
664+
665+ let null_buffer = array. keys ( ) . nulls ( ) . expect ( "HAS_NULL_KEYS implies nulls" ) ;
666+ let bit_chunks = null_buffer. inner ( ) . bit_chunks ( ) ;
667+ let chunk_len = bit_chunks. chunk_len ( ) ;
668+
669+ for ( chunk_idx, mask) in bit_chunks. iter ( ) . enumerate ( ) {
670+ let base = chunk_idx * 64 ;
671+ let keys_chunk = & keys[ base..base + 64 ] ;
672+ let hashes_chunk = & mut hashes_buffer[ base..base + 64 ] ;
673+ if mask == u64:: MAX {
674+ if HAS_NULL_VALUES {
675+ scatter_dict_dense_filter_values :: < K , MULTI_COL > (
676+ keys_chunk,
677+ dict_values. as_ref ( ) ,
678+ dict_hashes,
679+ hashes_chunk,
680+ ) ;
681+ } else {
682+ scatter_dict_dense :: < K , MULTI_COL > ( keys_chunk, dict_hashes, hashes_chunk) ;
550683 }
684+ } else if mask != 0 {
685+ scatter_dict_blend_chunk :: < K , HAS_NULL_VALUES , MULTI_COL > (
686+ keys_chunk,
687+ dict_values. as_ref ( ) ,
688+ dict_hashes,
689+ hashes_chunk,
690+ mask,
691+ ) ;
551692 }
552- } else {
553- for ( hash, key) in hashes_buffer. iter_mut ( ) . zip ( array. keys ( ) . values ( ) ) {
554- let idx = key. as_usize ( ) ;
693+ }
694+
695+ let remainder_len = bit_chunks. remainder_len ( ) ;
696+ if remainder_len > 0 {
697+ let remainder = bit_chunks. remainder_bits ( ) ;
698+ let base = chunk_len * 64 ;
699+ let mut remaining = remainder;
700+ while remaining != 0 {
701+ let i = remaining. trailing_zeros ( ) as usize ;
702+ remaining &= remaining - 1 ;
703+ if i >= remainder_len {
704+ break ;
705+ }
706+ let idx = keys[ base + i] . as_usize ( ) ;
555707 if !HAS_NULL_VALUES || dict_values. is_valid ( idx) {
556- if MULTI_COL {
557- * hash = combine_hashes ( dict_hashes[ idx] , * hash) ;
558- } else {
559- * hash = dict_hashes[ idx] ;
560- }
708+ let dh = unsafe { * dict_hashes. get_unchecked ( idx) } ;
709+ hashes_buffer[ base + i] =
710+ apply_dict_hash :: < MULTI_COL > ( hashes_buffer[ base + i] , dh) ;
561711 }
562712 }
563713 }
0 commit comments