@@ -30,7 +30,8 @@ use crate::{
3030 test:: TestMemoryExec ,
3131} ;
3232use arrow:: array:: {
33- BooleanArray , FixedSizeBinaryArray , Int64Array , StringArray , StringViewArray ,
33+ BooleanArray , FixedSizeBinaryArray , Int64Array , NullArray , StringArray ,
34+ StringViewArray , StructArray ,
3435} ;
3536use arrow:: buffer:: Buffer ;
3637use arrow:: compute:: kernels:: sort:: SortOptions ;
@@ -698,6 +699,125 @@ fn prepared_concat_admits_validity_for_non_nullable_input_arrays() -> Result<()>
698699 Ok ( ( ) )
699700}
700701
702+ #[ tokio:: test]
703+ async fn prepared_inlist_shares_accounted_build_buffers ( ) -> Result < ( ) > {
704+ let build = batch ( vec ! [ Some ( 1 ) , Some ( 1_000_000 ) , None , Some ( 1 ) ] ) ;
705+ for keys in [ 1 , 2 ] {
706+ let on = ( 0 ..keys)
707+ . map ( |index| {
708+ let name = build. schema ( ) . field ( index) . name ( ) . clone ( ) ;
709+ let key: PhysicalExprRef = Arc :: new ( Column :: new ( & name, index) ) ;
710+ ( Arc :: clone ( & key) , key)
711+ } )
712+ . collect ( ) ;
713+ let base = join ( build. schema ( ) , build. clone ( ) ) ?
714+ . builder ( )
715+ . with_on ( on)
716+ . build ( ) ?;
717+ for batches in [
718+ vec ! [ build. clone( ) ] ,
719+ vec ! [ build. slice( 0 , 2 ) , build. slice( 2 , 2 ) ] ,
720+ ] {
721+ let pool: Arc < dyn MemoryPool > = Arc :: new ( GreedyMemoryPool :: new ( 1 << 20 ) ) ;
722+ let prepared = prepare ( & base, batches. clone ( ) , Arc :: clone ( & pool) ) . await ?;
723+ let PushdownStrategy :: InList ( values) = & prepared. build . membership else {
724+ panic ! ( "expected IN-list membership" ) ;
725+ } ;
726+ let columns = prepared. build . batch . columns ( ) ;
727+ if keys == 1 {
728+ assert ! ( Arc :: ptr_eq( values, & columns[ 0 ] ) ) ;
729+ } else {
730+ let values = values. as_any ( ) . downcast_ref :: < StructArray > ( ) . unwrap ( ) ;
731+ for ( value, column) in values. columns ( ) . iter ( ) . zip ( columns) {
732+ assert ! ( Arc :: ptr_eq( value, column) ) ;
733+ }
734+ }
735+ let bytes = prepared. reserved_bytes ( ) ;
736+ let mut counter = RecordBatchMemoryCounter :: new ( ) ;
737+ let batch_bytes = counter. count_batch ( & prepared. build . batch ) ;
738+ assert ! ( bytes > batch_bytes, "hash storage must also be reserved" ) ;
739+ let membership_batch =
740+ RecordBatch :: try_from_iter ( [ ( "membership" , Arc :: clone ( values) ) ] ) ?;
741+ assert_eq ! ( counter. count_batch( & membership_batch) , 0 ) ;
742+ drop ( membership_batch) ;
743+
744+ // Disabling IN-list membership changes no retained key buffers.
745+ let mut config = ConfigOptions :: default ( ) ;
746+ config. optimizer . hash_join_inlist_pushdown_max_size = 0 ;
747+ let map = base
748+ . prepare_build (
749+ input ( batches, build. schema ( ) ) ,
750+ Arc :: clone ( & pool) ,
751+ Arc :: new ( config) ,
752+ )
753+ . await ?;
754+ assert ! ( matches!( map. build. membership, PushdownStrategy :: Map ( _) ) ) ;
755+ assert_eq ! ( map. reserved_bytes( ) , bytes) ;
756+ assert_eq ! ( pool. reserved( ) , 2 * bytes) ;
757+ drop ( map) ;
758+ let lease = Arc :: clone ( & prepared) ;
759+ drop ( prepared) ;
760+ assert_eq ! ( pool. reserved( ) , bytes) ;
761+ drop ( lease) ;
762+ assert_eq ! ( pool. reserved( ) , 0 ) ;
763+ }
764+ }
765+ Ok ( ( ) )
766+ }
767+
768+ #[ test]
769+ fn prepared_concat_rejects_under_admission_without_shrinking ( ) -> Result < ( ) > {
770+ let pool: Arc < dyn MemoryPool > = Arc :: new ( GreedyMemoryPool :: new ( 1024 ) ) ;
771+ let reservation = MemoryConsumer :: new ( "prepared concat test" ) . register ( & pool) ;
772+ // 128 bytes of map storage plus 128 bytes of input/copy allowance.
773+ reservation. try_grow ( 256 ) ?;
774+ let error =
775+ reconcile_prepared_copy_reservation ( & reservation, 64 , 64 , 129 ) . unwrap_err ( ) ;
776+ assert ! ( matches!( error, DataFusionError :: Internal ( _) ) ) ;
777+ assert ! (
778+ error
779+ . to_string( )
780+ . contains( "exceeded its admitted copy bound" )
781+ ) ;
782+ assert_eq ! ( reservation. size( ) , 256 ) ;
783+ assert ! ( reconcile_prepared_copy_reservation( & reservation, usize :: MAX , 1 , 0 ) . is_err( ) ) ;
784+ assert_eq ! ( reservation. size( ) , 256 ) ;
785+ reconcile_prepared_copy_reservation ( & reservation, 64 , 64 , 128 ) ?;
786+ assert_eq ! ( reservation. size( ) , 256 ) ;
787+ reconcile_prepared_copy_reservation ( & reservation, 64 , 64 , 80 ) ?;
788+ assert_eq ! ( reservation. size( ) , 208 ) ;
789+ drop ( reservation) ;
790+ assert_eq ! ( pool. reserved( ) , 0 ) ;
791+ Ok ( ( ) )
792+ }
793+
794+ #[ test]
795+ fn prepared_sizing_rejects_overflow ( ) -> Result < ( ) > {
796+ use NullEquality :: { NullEqualsNothing , NullEqualsNull } ;
797+ for ( rows, keys, equality) in [
798+ ( usize:: MAX / 8 + 1 , 1 , NullEqualsNull ) ,
799+ ( 0 , usize:: MAX , NullEqualsNothing ) ,
800+ ( 0 , usize:: MAX / 64 , NullEqualsNothing ) ,
801+ ( usize:: MAX / 8 , 0 , NullEqualsNothing ) ,
802+ ] {
803+ let error = prepared_scratch_bytes ( rows, keys, equality) . unwrap_err ( ) ;
804+ assert ! ( error. to_string( ) . contains( "scratch size overflow" ) ) ;
805+ }
806+ assert ! ( prepared_scratch_bytes( usize :: MAX / 8 , 1 , NullEqualsNull ) . is_ok( ) ) ;
807+
808+ // Null arrays can describe huge row counts without allocating data buffers.
809+ // Exercise both per-batch validity multiplication and cumulative addition.
810+ let array: ArrayRef = Arc :: new ( NullArray :: new ( usize:: MAX ) ) ;
811+ for ( columns, batches) in [ ( 8 , 1 ) , ( 4 , 2 ) ] {
812+ let batch = RecordBatch :: try_from_iter (
813+ ( 0 ..columns) . map ( |i| ( format ! ( "key{i}" ) , Arc :: clone ( & array) ) ) ,
814+ ) ?;
815+ let error = prepared_copy_bytes ( & vec ! [ batch; batches] ) . unwrap_err ( ) ;
816+ assert ! ( error. to_string( ) . contains( "copy size overflow" ) ) ;
817+ }
818+ Ok ( ( ) )
819+ }
820+
701821#[ tokio:: test]
702822async fn prepared_byte_keys_use_hash_membership ( ) -> Result < ( ) > {
703823 for array in [
0 commit comments