2121
2222use std:: fmt:: { Debug , Display , Formatter } ;
2323use std:: pin:: Pin ;
24- use std:: sync:: Arc ;
2524use std:: sync:: atomic:: { AtomicUsize , Ordering as AtomicOrdering } ;
25+ use std:: sync:: { Arc , Weak } ;
2626use std:: task:: { Context , Poll } ;
2727use std:: vec;
2828
@@ -33,7 +33,9 @@ use super::{
3333} ;
3434use crate :: coalesce:: LimitedBatchCoalescer ;
3535use crate :: execution_plan:: { CardinalityEffect , EvaluationType , SchedulingType } ;
36- use crate :: hash_utils:: create_hashes;
36+ #[ cfg( not( feature = "force_hash_collisions" ) ) ]
37+ use crate :: hash_utils:: scatter_dictionary_hashes;
38+ use crate :: hash_utils:: { create_hashes, hash_single_array} ;
3739use crate :: metrics:: { BaselineMetrics , SpillMetrics } ;
3840use crate :: projection:: { ProjectionExec , all_columns, make_with_child, update_expr} ;
3941use crate :: sorts:: streaming_merge:: StreamingMergeBuilder ;
@@ -46,9 +48,12 @@ use crate::{
4648 PlanProperties , ReplaceChildrenOptions , Statistics , validate_child_count,
4749} ;
4850
49- use arrow:: array:: { PrimitiveArray , RecordBatch , RecordBatchOptions , UInt64Array } ;
51+ use arrow:: array:: {
52+ Array , ArrayRef , PrimitiveArray , RecordBatch , RecordBatchOptions , UInt64Array ,
53+ } ;
5054use arrow:: compute:: take_arrays;
5155use arrow:: datatypes:: { DataType , Schema , SchemaRef , UInt32Type } ;
56+ use arrow:: downcast_dictionary_array;
5257use arrow_schema:: SortOptions ;
5358use datafusion_common:: ScalarValue ;
5459use datafusion_common:: config:: ConfigOptions ;
@@ -657,6 +662,7 @@ enum BatchPartitionerState {
657662 partition_reducer : StrengthReducedU64 ,
658663 hash_buffer : Vec < u64 > ,
659664 indices : Vec < Vec < u32 > > ,
665+ dictionary_hash_cache : Vec < DictionaryHashCache > ,
660666 } ,
661667 RoundRobin {
662668 num_partitions : usize ,
@@ -672,6 +678,40 @@ enum BatchPartitionerState {
672678 } ,
673679}
674680
681+ #[ derive( Default ) ]
682+ struct DictionaryHashCache {
683+ /// Dictionary values these hashes belong to.
684+ values : Option < Weak < dyn Array > > ,
685+ /// Hash for each dictionary value.
686+ hashes : Vec < u64 > ,
687+ /// Number of times values were hashed.
688+ #[ cfg( test) ]
689+ rebuilds : usize ,
690+ }
691+
692+ impl DictionaryHashCache {
693+ fn update ( & mut self , values : & ArrayRef ) -> Result < & [ u64 ] > {
694+ let unchanged = self . values . as_ref ( ) . is_some_and ( |cached| {
695+ std:: ptr:: addr_eq ( cached. as_ptr ( ) , Arc :: as_ptr ( values) )
696+ } ) ;
697+ if !unchanged {
698+ self . hashes . clear ( ) ;
699+ self . hashes . resize ( values. len ( ) , 0 ) ;
700+ create_hashes (
701+ [ values] ,
702+ REPARTITION_RANDOM_STATE . random_state ( ) ,
703+ & mut self . hashes ,
704+ ) ?;
705+ self . values = Some ( Arc :: downgrade ( values) ) ;
706+ #[ cfg( test) ]
707+ {
708+ self . rebuilds += 1 ;
709+ }
710+ }
711+ Ok ( & self . hashes )
712+ }
713+ }
714+
675715/// Fixed RandomState used for hash repartitioning to ensure consistent behavior across
676716/// executions and runs.
677717pub const REPARTITION_RANDOM_STATE : SeededRandomState = SeededRandomState :: with_seed ( 0 ) ;
@@ -1036,6 +1076,7 @@ impl BatchPartitioner {
10361076 partition_reducer : StrengthReducedU64 :: new ( num_partitions as u64 ) ,
10371077 hash_buffer : vec ! [ ] ,
10381078 indices : vec ! [ vec![ ] ; num_partitions] ,
1079+ dictionary_hash_cache : vec ! [ ] ,
10391080 } ,
10401081 timer,
10411082 } )
@@ -1250,6 +1291,7 @@ impl BatchPartitioner {
12501291 partition_reducer,
12511292 hash_buffer,
12521293 indices,
1294+ dictionary_hash_cache,
12531295 } => {
12541296 // Tracking time required for distributing indexes across output partitions
12551297 let timer = self . timer . timer ( ) ;
@@ -1260,11 +1302,55 @@ impl BatchPartitioner {
12601302 hash_buffer. clear ( ) ;
12611303 hash_buffer. resize ( batch. num_rows ( ) , 0 ) ;
12621304
1263- create_hashes (
1264- & arrays,
1265- REPARTITION_RANDOM_STATE . random_state ( ) ,
1266- hash_buffer,
1267- ) ?;
1305+ dictionary_hash_cache
1306+ . resize_with ( arrays. len ( ) , DictionaryHashCache :: default) ;
1307+ for ( index, ( array, cache) ) in arrays
1308+ . iter ( )
1309+ . zip ( dictionary_hash_cache. iter_mut ( ) )
1310+ . enumerate ( )
1311+ {
1312+ let array = array. as_ref ( ) ;
1313+ if matches ! ( array. data_type( ) , DataType :: Dictionary ( _, _) ) {
1314+ downcast_dictionary_array ! {
1315+ array => {
1316+ let hashes = cache. update( array. values( ) ) ?;
1317+ #[ cfg( not( feature = "force_hash_collisions" ) ) ]
1318+ scatter_dictionary_hashes(
1319+ array,
1320+ hashes,
1321+ hash_buffer,
1322+ index != 0 ,
1323+ ) ;
1324+ #[ cfg( feature = "force_hash_collisions" ) ]
1325+ let _ = hashes;
1326+ #[ cfg( feature = "force_hash_collisions" ) ]
1327+ hash_single_array(
1328+ array,
1329+ REPARTITION_RANDOM_STATE . random_state( ) ,
1330+ hash_buffer,
1331+ index != 0 ,
1332+ ) ?;
1333+ } ,
1334+ _ => unreachable!( ) ,
1335+ }
1336+ } else {
1337+ * cache = DictionaryHashCache :: default ( ) ;
1338+ hash_single_array (
1339+ array,
1340+ REPARTITION_RANDOM_STATE . random_state ( ) ,
1341+ hash_buffer,
1342+ index != 0 ,
1343+ ) ?;
1344+ }
1345+ }
1346+
1347+ if arrays. is_empty ( ) {
1348+ create_hashes (
1349+ arrays,
1350+ REPARTITION_RANDOM_STATE . random_state ( ) ,
1351+ hash_buffer,
1352+ ) ?;
1353+ }
12681354
12691355 for v in indices. iter_mut ( ) {
12701356 v. clear ( ) ;
@@ -2660,10 +2746,10 @@ mod tests {
26602746 } ;
26612747
26622748 use arrow:: array:: {
2663- Array , ArrayRef , Decimal128Array , Int64Array , StringArray ,
2664- TimestampNanosecondArray , UInt32Array ,
2749+ Array , ArrayRef , Decimal128Array , DictionaryArray , Int32Array , Int64Array ,
2750+ StringArray , TimestampNanosecondArray , UInt32Array ,
26652751 } ;
2666- use arrow:: datatypes:: { DataType , Field , Schema , TimeUnit } ;
2752+ use arrow:: datatypes:: { DataType , Field , Int32Type , Schema , TimeUnit } ;
26672753 use datafusion_common:: ScalarValue ;
26682754 use datafusion_common:: cast:: { as_string_array, as_uint32_array} ;
26692755 use datafusion_common:: config:: SpillCompression ;
@@ -2918,6 +3004,83 @@ mod tests {
29183004 Ok ( ( ) )
29193005 }
29203006
3007+ #[ test]
3008+ fn dictionary_hash_cache_rebuilds_when_values_change ( ) -> Result < ( ) > {
3009+ let first: ArrayRef =
3010+ Arc :: new ( StringArray :: from ( vec ! [ Some ( "a" ) , None , Some ( "ccc" ) ] ) ) ;
3011+ let second: ArrayRef =
3012+ Arc :: new ( StringArray :: from ( vec ! [ Some ( "x" ) , None , Some ( "zzz" ) ] ) ) ;
3013+ let mut cache = DictionaryHashCache :: default ( ) ;
3014+
3015+ let first_hashes = cache. update ( & first) ?. to_vec ( ) ;
3016+ cache. update ( & first) ?;
3017+ assert_eq ! ( cache. rebuilds, 1 ) ;
3018+
3019+ let second_hashes = cache. update ( & second) ?. to_vec ( ) ;
3020+ assert_eq ! ( cache. rebuilds, 2 ) ;
3021+ assert_ne ! ( first_hashes, second_hashes) ;
3022+ Ok ( ( ) )
3023+ }
3024+
3025+ #[ tokio:: test]
3026+ async fn hash_partition_reuses_dictionary_values_across_batches ( ) -> Result < ( ) > {
3027+ let dictionary_type =
3028+ DataType :: Dictionary ( Box :: new ( DataType :: Int32 ) , Box :: new ( DataType :: Utf8 ) ) ;
3029+ let schema = Arc :: new ( Schema :: new ( vec ! [
3030+ Field :: new( "key" , dictionary_type, true ) ,
3031+ Field :: new( "id" , DataType :: UInt32 , false ) ,
3032+ ] ) ) ;
3033+ let values: ArrayRef = Arc :: new ( StringArray :: from ( vec ! [ "a" , "bb" , "ccc" ] ) ) ;
3034+ let make_batch = |keys : Vec < Option < i32 > > , ids : Vec < u32 > | -> Result < RecordBatch > {
3035+ let dictionary: ArrayRef = Arc :: new ( DictionaryArray :: < Int32Type > :: try_new (
3036+ Int32Array :: from ( keys) ,
3037+ Arc :: clone ( & values) ,
3038+ ) ?) ;
3039+ Ok ( RecordBatch :: try_new (
3040+ Arc :: clone ( & schema) ,
3041+ vec ! [ dictionary, Arc :: new( UInt32Array :: from( ids) ) ] ,
3042+ ) ?)
3043+ } ;
3044+ let batches = vec ! [
3045+ make_batch( vec![ Some ( 0 ) , Some ( 1 ) , Some ( 2 ) , None ] , vec![ 0 , 1 , 2 , 3 ] ) ?,
3046+ make_batch( vec![ Some ( 2 ) , Some ( 0 ) , None , Some ( 1 ) ] , vec![ 4 , 5 , 6 , 7 ] ) ?,
3047+ ] ;
3048+
3049+ let mut expected = vec ! [ Vec :: new( ) ; 4 ] ;
3050+ for batch in & batches {
3051+ let mut hashes = vec ! [ 0 ; batch. num_rows( ) ] ;
3052+ create_hashes (
3053+ [ batch. column ( 0 ) ] ,
3054+ REPARTITION_RANDOM_STATE . random_state ( ) ,
3055+ & mut hashes,
3056+ ) ?;
3057+ let ids = as_uint32_array ( batch. column ( 1 ) ) ?;
3058+ for ( hash, id) in hashes. into_iter ( ) . zip ( ids. values ( ) . iter ( ) . copied ( ) ) {
3059+ expected[ ( hash % 4 ) as usize ] . push ( id) ;
3060+ }
3061+ }
3062+
3063+ let output = repartition (
3064+ & schema,
3065+ vec ! [ batches] ,
3066+ Partitioning :: Hash ( vec ! [ col( "key" , & schema) ?] , 4 ) ,
3067+ )
3068+ . await ?;
3069+ let mut actual = vec ! [ Vec :: new( ) ; 4 ] ;
3070+ for ( partition, batches) in output. iter ( ) . enumerate ( ) {
3071+ for batch in batches {
3072+ actual[ partition]
3073+ . extend ( as_uint32_array ( batch. column ( 1 ) ) ?. values ( ) . iter ( ) . copied ( ) ) ;
3074+ }
3075+ }
3076+ for values in expected. iter_mut ( ) . chain ( actual. iter_mut ( ) ) {
3077+ values. sort_unstable ( ) ;
3078+ }
3079+
3080+ assert_eq ! ( actual, expected) ;
3081+ Ok ( ( ) )
3082+ }
3083+
29213084 #[ tokio:: test]
29223085 async fn many_to_many_range_partition ( ) -> Result < ( ) > {
29233086 let schema = test_schema ( false ) ;
0 commit comments