77use std:: { fmt:: Display , hash:: Hash , sync:: Arc } ;
88
99use arrow:: {
10- array:: { ArrayRef , BooleanArray , BooleanBufferBuilder , UInt64Array } ,
10+ array:: { ArrayRef , BooleanArray , BooleanBufferBuilder , UInt32Array , UInt64Array } ,
11+ compute:: take,
1112 datatypes:: { DataType , Schema } ,
1213 record_batch:: RecordBatch ,
1314} ;
@@ -352,8 +353,7 @@ impl HashTableLookupExpr {
352353 /// for pruning only, or `None`
353354 ///
354355 /// A single non-empty map is probed directly without routing, so it may be
355- /// any [`Map`] and `router` is ignored. Otherwise `router` is required and
356- /// every map must be a [`Map::HashMap`].
356+ /// any [`Map`] and `router` is ignored. Otherwise `router` is required.
357357 ///
358358 /// # Public Only for Internal Use:
359359 /// `datafusion-proto` tests require this constructor, but it is not part of
@@ -375,21 +375,10 @@ impl HashTableLookupExpr {
375375 "a lookup over anything but one non-empty map needs a partition router"
376376 ) ;
377377 }
378- ( Some ( router) , _) => {
379- if maps
380- . iter ( )
381- . flatten ( )
382- . any ( |map| !matches ! ( map. as_ref( ) , Map :: HashMap ( _) ) )
383- {
384- return internal_err ! (
385- "routed lookups only support hash table partitions"
386- ) ;
387- }
388- Lookup :: Routed {
389- router,
390- maps : maps. into ( ) ,
391- }
392- }
378+ ( Some ( router) , _) => Lookup :: Routed {
379+ router,
380+ maps : maps. into ( ) ,
381+ } ,
393382 } ;
394383 Ok ( Self {
395384 on_columns,
@@ -441,16 +430,23 @@ impl HashTableLookupExpr {
441430 for ( partition, map) in maps. iter ( ) . enumerate ( ) {
442431 // An empty partition holds no key, so its rows stay unmatched.
443432 let Some ( map) = map else { continue } ;
444- let Map :: HashMap ( map) = map. as_ref ( ) else {
445- return internal_err ! (
446- "routed lookups only support hash table partitions"
447- ) ;
448- } ;
449433 let run = starts[ partition] ..starts[ partition + 1 ] ;
450434 if run. is_empty ( ) {
451435 continue ;
452436 }
453- let hits = map. contain_hashes ( & sorted_hashes[ run. clone ( ) ] ) ;
437+ let hits = match map. as_ref ( ) {
438+ Map :: HashMap ( map) => map. contain_hashes ( & sorted_hashes[ run. clone ( ) ] ) ,
439+ // An array map is keyed by value, not hash, so it probes
440+ // the run's own key values.
441+ Map :: ArrayMap ( map) => {
442+ let indices = UInt32Array :: from ( rows[ run. clone ( ) ] . to_vec ( ) ) ;
443+ let run_keys = keys
444+ . iter ( )
445+ . map ( |key| take ( key. as_ref ( ) , & indices, None ) )
446+ . collect :: < std:: result:: Result < Vec < _ > , _ > > ( ) ?;
447+ map. contain_keys ( & run_keys) ?
448+ }
449+ } ;
454450 for hit in hits. values ( ) . set_indices ( ) {
455451 matched. set_bit ( rows[ run. start + hit] as usize , true ) ;
456452 }
@@ -639,7 +635,9 @@ fn evaluate_columns(
639635#[ cfg( test) ]
640636mod tests {
641637 use super :: * ;
638+ use crate :: joins:: array_map:: ArrayMap ;
642639 use crate :: joins:: join_hash_map:: JoinHashMapU32 ;
640+ use arrow:: array:: AsArray ;
643641 use datafusion_common:: ScalarValue ;
644642 use datafusion_physical_expr:: expressions:: Column ;
645643 use std:: collections:: hash_map:: DefaultHasher ;
@@ -1111,7 +1109,7 @@ mod tests {
11111109 ) ] ) ) ;
11121110 let batch = RecordBatch :: try_new ( schema, vec ! [ Arc :: clone( probe) ] ) ?;
11131111 let matched = expr. evaluate ( & batch) ?. into_array ( batch. num_rows ( ) ) ?;
1114- Ok ( arrow :: array :: AsArray :: as_boolean ( & matched ) . clone ( ) )
1112+ Ok ( matched . as_boolean ( ) . clone ( ) )
11151113 }
11161114
11171115 /// A hash table of `keys`, hashed with `random_state` as the build side
@@ -1125,6 +1123,18 @@ mod tests {
11251123 Ok ( Arc :: new ( Map :: HashMap ( Box :: new ( map) ) ) )
11261124 }
11271125
1126+ /// An array map of `keys`, as a dense build side gets.
1127+ fn array_map ( keys : & ArrayRef ) -> Result < Arc < Map > > {
1128+ let keys = keys. as_primitive :: < arrow:: datatypes:: Int32Type > ( ) ;
1129+ let key = |value : Option < i32 > | {
1130+ ArrayMap :: key_to_u64 ( & ScalarValue :: Int32 ( value) ) . expect ( "Int32 key" )
1131+ } ;
1132+ let min = key ( arrow:: compute:: min ( keys) ) ;
1133+ let max = key ( arrow:: compute:: max ( keys) ) ;
1134+ let keys: ArrayRef = Arc :: new ( keys. clone ( ) ) ;
1135+ Ok ( Arc :: new ( Map :: ArrayMap ( ArrayMap :: try_new ( & keys, min, max) ?) ) )
1136+ }
1137+
11281138 /// One randomly generated scenario, see [`routed_lookup_fuzz`].
11291139 fn run_routed_lookup_scenario ( seed : u64 ) -> Result < ( ) > {
11301140 use crate :: repartition:: { BatchPartitioner , REPARTITION_RANDOM_STATE } ;
@@ -1166,7 +1176,12 @@ mod tests {
11661176 . partition (
11671177 RecordBatch :: try_new ( schema, vec ! [ Arc :: clone( & build) ] ) ?,
11681178 |partition, batch| {
1169- maps[ partition] = Some ( hash_table ( batch. column ( 0 ) , & random_state) ?) ;
1179+ // A dense partition builds an array map, others a hash table.
1180+ maps[ partition] = Some ( if rng. random_bool ( 0.5 ) {
1181+ array_map ( batch. column ( 0 ) ) ?
1182+ } else {
1183+ hash_table ( batch. column ( 0 ) , & random_state) ?
1184+ } ) ;
11701185 Ok ( ( ) )
11711186 } ,
11721187 ) ?;
0 commit comments