@@ -66,7 +66,7 @@ use datafusion_execution::memory_pool::{
6666use datafusion_execution:: runtime_env:: RuntimeEnv ;
6767use datafusion_physical_expr:: LexOrdering ;
6868use datafusion_physical_expr:: PhysicalExpr ;
69- use datafusion_physical_expr:: expressions:: { DynamicFilterPhysicalExpr , lit} ;
69+ use datafusion_physical_expr:: expressions:: { Column , DynamicFilterPhysicalExpr , lit} ;
7070
7171use futures:: { Stream , StreamExt , TryStreamExt } ;
7272use log:: { debug, trace} ;
@@ -784,13 +784,13 @@ impl ExternalSorter {
784784 let ( sizes, shared_input_size) = self . sort_run_sizes ( & runs) ?;
785785 let shared_input =
786786 shared_input_size. map ( |size| Arc :: new ( self . reservation . split ( size) ) ) ;
787- let batch_size = self . merge_run_batch_size ( ) ;
788787 let jobs =
789788 runs. into_iter ( )
790789 . zip ( sizes)
791790 . enumerate ( )
792791 . map ( move |( index, ( batch, size) ) | {
793792 let reservation = self . reservation . split ( size) ;
793+ let batch_size = self . merge_run_batch_size ( & batch) ?;
794794 let sort = self . sort_run (
795795 batch,
796796 reservation,
@@ -978,11 +978,11 @@ impl ExternalSorter {
978978 let ( sizes, shared_input_size) = self . sort_run_sizes ( & runs) ?;
979979 let shared_input =
980980 shared_input_size. map ( |size| Arc :: new ( self . reservation . split ( size) ) ) ;
981- let batch_size = self . merge_run_batch_size ( ) ;
982981 runs. into_iter ( )
983982 . zip ( sizes)
984983 . map ( |( run, size) | {
985984 let reservation = self . reservation . split ( size) ;
985+ let batch_size = self . merge_run_batch_size ( & run) ?;
986986 self . sort_batch_stream ( run, reservation, shared_input. clone ( ) , batch_size)
987987 } )
988988 . collect ( )
@@ -1038,12 +1038,6 @@ impl ExternalSorter {
10381038 let expressions = self . expr . clone ( ) ;
10391039 let merge_pool = Arc :: clone ( & self . merge_pool ) ;
10401040 let elapsed_compute = self . metrics . baseline . elapsed_compute ( ) . clone ( ) ;
1041- // Split batches of views sharing bytes would each be charged for them
1042- let batch_size = if batch_size < self . batch_size && views_share_bytes ( & batch) {
1043- self . batch_size
1044- } else {
1045- batch_size
1046- } ;
10471041 // Set when `merge_run_batch_size` split this run
10481042 let split = batch_size < self . batch_size ;
10491043
@@ -1096,15 +1090,17 @@ impl ExternalSorter {
10961090 } )
10971091 }
10981092
1099- /// Batch size of the sorted runs that feed a merge. A merge in the row format
1100- /// holds the keys of every run's current batch at once, and these keys can be
1101- /// larger than the batches. With a memory limit, runs therefore use an eighth
1102- /// of `batch_size`, but at least 1024 rows to keep the per-batch merge cost
1103- /// small. The merge output still uses `batch_size`. Runs keep `batch_size`
1104- /// when their batches would share buffers other than the data of top-level
1105- /// view columns, such as dictionary values, since every batch is charged for
1106- /// them. `sort_run` also keeps it for a run whose views share bytes.
1107- fn merge_run_batch_size ( & self ) -> usize {
1093+ /// Batch size of the sorted run `run` when it feeds a merge. A merge in the
1094+ /// row format holds the keys of every run's current batch at once. The sort
1095+ /// reserves about twice the run, so the keys fit next to the sorted copy if
1096+ /// they are at most as large as the run. Otherwise, with a memory limit, the
1097+ /// run uses an eighth of `batch_size`, but at least 1024 rows to keep the
1098+ /// per-batch merge cost small. The merge output still uses `batch_size`.
1099+ ///
1100+ /// Runs keep `batch_size` when their batches would share buffers other than
1101+ /// the data of top-level view columns, such as dictionary values, or when the
1102+ /// views of `run` share bytes, since every batch is charged for them.
1103+ fn merge_run_batch_size ( & self , run : & RecordBatch ) -> Result < usize > {
11081104 let limited = !matches ! (
11091105 self . runtime. memory_pool. memory_limit( ) ,
11101106 MemoryLimit :: Infinite
@@ -1114,11 +1110,34 @@ impl ExternalSorter {
11141110 sorted_copy_keeps_buffers ( data_type)
11151111 && !matches ! ( data_type, DataType :: Utf8View | DataType :: BinaryView )
11161112 } ) ;
1117- if limited && !shares_buffers && merges_in_row_format ( & self . expr , & self . schema ) {
1118- ( self . batch_size / 8 ) . max ( 1024 ) . min ( self . batch_size )
1113+ if limited
1114+ && !shares_buffers
1115+ && merges_in_row_format ( & self . expr , & self . schema )
1116+ && self . encoded_keys_may_exceed ( run) ?
1117+ && !views_share_bytes ( run)
1118+ {
1119+ Ok ( ( self . batch_size / 8 ) . max ( 1024 ) . min ( self . batch_size ) )
11191120 } else {
1120- self . batch_size
1121+ Ok ( self . batch_size )
1122+ }
1123+ }
1124+
1125+ /// Returns whether the row-format keys of `run` may be larger than `run`. A
1126+ /// computed key counts as larger, since sizing it would evaluate it again.
1127+ fn encoded_keys_may_exceed ( & self , run : & RecordBatch ) -> Result < bool > {
1128+ // An 8-byte offset per row, then a null byte and the bytes of each key
1129+ let mut encoded = run. num_rows ( ) * size_of :: < usize > ( ) ;
1130+ for sort_expr in self . expr . iter ( ) {
1131+ let Some ( column) = sort_expr. expr . downcast_ref :: < Column > ( ) else {
1132+ return Ok ( true ) ;
1133+ } ;
1134+ encoded += run. num_rows ( )
1135+ + run
1136+ . column ( column. index ( ) )
1137+ . to_data ( )
1138+ . get_slice_memory_size ( ) ?;
11211139 }
1140+ Ok ( encoded > run. get_sliced_size ( ) ?)
11221141 }
11231142
11241143 /// If this sort may spill, pre-allocates
0 commit comments