@@ -46,10 +46,14 @@ use crate::{
4646 PlanProperties , ReplaceChildrenOptions , Statistics , validate_child_count,
4747} ;
4848
49- use arrow:: array:: { PrimitiveArray , RecordBatch , RecordBatchOptions , UInt64Array } ;
49+ use arrow:: array:: {
50+ Array , ArrayRef , AsArray , PrimitiveArray , RecordBatch , RecordBatchOptions ,
51+ UInt64Array ,
52+ } ;
5053use arrow:: compute:: take_arrays;
5154use arrow:: datatypes:: { DataType , Schema , SchemaRef , UInt32Type } ;
5255use arrow_schema:: SortOptions ;
56+ use arrow_select:: dictionary:: garbage_collect_any_dictionary;
5357use datafusion_common:: ScalarValue ;
5458use datafusion_common:: config:: ConfigOptions ;
5559use datafusion_common:: stats:: Precision ;
@@ -1382,7 +1386,8 @@ impl BatchPartitioner {
13821386 partition_ranges
13831387 . into_iter ( )
13841388 . map ( |( partition, start, len) | {
1385- Ok ( ( partition, reordered_batch. slice ( start, len) ) )
1389+ let batch = reordered_batch. slice ( start, len) ;
1390+ Ok ( ( partition, compact_sparse_dictionaries ( batch) ?) )
13861391 } )
13871392 . collect ( )
13881393 } ;
@@ -1391,6 +1396,44 @@ impl BatchPartitioner {
13911396 }
13921397}
13931398
1399+ /// Remove unreferenced values from dictionaries made sparse by repartitioning.
1400+ ///
1401+ /// `take_arrays` reorders dictionary keys but preserves their complete values
1402+ /// arrays. Each output partition then references only a fraction of those
1403+ /// values, and coalescing the partitions would otherwise hash every retained
1404+ /// value while merging their dictionaries.
1405+ fn compact_sparse_dictionaries ( batch : RecordBatch ) -> Result < RecordBatch > {
1406+ let threshold = batch. num_rows ( ) . saturating_mul ( 4 ) ;
1407+ let mut changed = false ;
1408+ let columns = batch
1409+ . columns ( )
1410+ . iter ( )
1411+ . map ( |column| {
1412+ let DataType :: Dictionary ( _, _) = column. data_type ( ) else {
1413+ return Ok ( Arc :: clone ( column) ) ;
1414+ } ;
1415+ let dictionary = column. as_any_dictionary ( ) ;
1416+ if dictionary. values ( ) . len ( ) <= threshold {
1417+ return Ok ( Arc :: clone ( column) ) ;
1418+ }
1419+ changed = true ;
1420+ Ok ( garbage_collect_any_dictionary ( dictionary) ?)
1421+ } )
1422+ . collect :: < Result < Vec < ArrayRef > > > ( ) ?;
1423+
1424+ if !changed {
1425+ return Ok ( batch) ;
1426+ }
1427+
1428+ let mut options = RecordBatchOptions :: new ( ) ;
1429+ options = options. with_row_count ( Some ( batch. num_rows ( ) ) ) ;
1430+ Ok ( RecordBatch :: try_new_with_options (
1431+ batch. schema ( ) ,
1432+ columns,
1433+ & options,
1434+ ) ?)
1435+ }
1436+
13941437/// Maps `N` input partitions to `M` output partitions based on a
13951438/// [`Partitioning`] scheme.
13961439///
@@ -2660,10 +2703,10 @@ mod tests {
26602703 } ;
26612704
26622705 use arrow:: array:: {
2663- Array , ArrayRef , Decimal128Array , Int64Array , StringArray ,
2664- TimestampNanosecondArray , UInt32Array ,
2706+ Array , ArrayRef , Decimal128Array , DictionaryArray , Int32Array , Int64Array ,
2707+ StringArray , TimestampNanosecondArray , UInt32Array ,
26652708 } ;
2666- use arrow:: datatypes:: { DataType , Field , Schema , TimeUnit } ;
2709+ use arrow:: datatypes:: { DataType , Field , Int32Type , Schema , TimeUnit } ;
26672710 use datafusion_common:: ScalarValue ;
26682711 use datafusion_common:: cast:: { as_string_array, as_uint32_array} ;
26692712 use datafusion_common:: config:: SpillCompression ;
@@ -2677,6 +2720,46 @@ mod tests {
26772720 } ;
26782721 use insta:: assert_snapshot;
26792722
2723+ #[ test]
2724+ fn compact_sparse_dictionaries_only_rebuilds_sparse_columns ( ) -> Result < ( ) > {
2725+ let values = Arc :: new ( StringArray :: from_iter_values (
2726+ ( 0 ..32 ) . map ( |index| format ! ( "value-{index}" ) ) ,
2727+ ) ) ;
2728+ let sparse = Arc :: new ( DictionaryArray :: < Int32Type > :: new (
2729+ Int32Array :: from ( vec ! [ 31 , 2 , 7 , 2 ] ) ,
2730+ values,
2731+ ) ) as ArrayRef ;
2732+ let schema = Arc :: new ( Schema :: new ( vec ! [ Field :: new(
2733+ "dictionary" ,
2734+ sparse. data_type( ) . clone( ) ,
2735+ false ,
2736+ ) ] ) ) ;
2737+ let compacted = compact_sparse_dictionaries ( RecordBatch :: try_new (
2738+ Arc :: clone ( & schema) ,
2739+ vec ! [ Arc :: clone( & sparse) ] ,
2740+ ) ?) ?;
2741+ let dictionary = compacted. column ( 0 ) . as_dictionary :: < Int32Type > ( ) ;
2742+ assert_eq ! ( dictionary. values( ) . len( ) , 3 ) ;
2743+ let compacted_values = dictionary. values ( ) . as_string :: < i32 > ( ) ;
2744+ assert_eq ! (
2745+ ( 0 ..dictionary. len( ) )
2746+ . map( |row| compacted_values. value( dictionary. key( row) . unwrap( ) ) )
2747+ . collect:: <Vec <_>>( ) ,
2748+ vec![ "value-31" , "value-2" , "value-7" , "value-2" ]
2749+ ) ;
2750+
2751+ let dense_values = Arc :: new ( StringArray :: from ( vec ! [ "a" , "b" , "c" , "d" ] ) ) ;
2752+ let dense = Arc :: new ( DictionaryArray :: < Int32Type > :: new (
2753+ Int32Array :: from ( vec ! [ 0 , 1 , 2 , 3 ] ) ,
2754+ dense_values,
2755+ ) ) as ArrayRef ;
2756+ let dense_batch = RecordBatch :: try_new ( schema, vec ! [ Arc :: clone( & dense) ] ) ?;
2757+ let dense_output = compact_sparse_dictionaries ( dense_batch) ?;
2758+ assert ! ( Arc :: ptr_eq( dense_output. column( 0 ) , & dense) ) ;
2759+
2760+ Ok ( ( ) )
2761+ }
2762+
26802763 #[ derive( Debug ) ]
26812764 struct UnboundedTestPartition {
26822765 schema : SchemaRef ,
0 commit comments