@@ -3497,22 +3497,55 @@ async fn test_passthrough_wrapper_projection_keeps_ordering() -> Result<()> {
34973497
34983498#[ tokio:: test]
34993499async fn test_unique_build_hash_join_unbounded_limit ( ) -> Result < ( ) > {
3500- check_unique_build_join_unbounded_limit ( JoinType :: Inner , false ) . await
3500+ check_unique_build_join_unbounded_limit (
3501+ JoinType :: Inner ,
3502+ false ,
3503+ NullEquality :: NullEqualsNothing ,
3504+ )
3505+ . await
35013506}
35023507
35033508#[ tokio:: test]
35043509async fn test_unique_build_right_join_unbounded_limit ( ) -> Result < ( ) > {
3505- check_unique_build_join_unbounded_limit ( JoinType :: Right , false ) . await
3510+ check_unique_build_join_unbounded_limit (
3511+ JoinType :: Right ,
3512+ false ,
3513+ NullEquality :: NullEqualsNothing ,
3514+ )
3515+ . await
35063516}
35073517
35083518#[ tokio:: test]
35093519async fn test_unique_build_left_sort_merge_join_unbounded_limit ( ) -> Result < ( ) > {
3510- check_unique_build_join_unbounded_limit ( JoinType :: Left , false ) . await
3520+ check_unique_build_join_unbounded_limit (
3521+ JoinType :: Left ,
3522+ false ,
3523+ NullEquality :: NullEqualsNothing ,
3524+ )
3525+ . await
35113526}
35123527
35133528#[ tokio:: test]
35143529async fn test_unique_build_computed_key_join_unbounded_limit ( ) -> Result < ( ) > {
3515- check_unique_build_join_unbounded_limit ( JoinType :: Inner , true ) . await
3530+ check_unique_build_join_unbounded_limit (
3531+ JoinType :: Inner ,
3532+ true ,
3533+ NullEquality :: NullEqualsNothing ,
3534+ )
3535+ . await
3536+ }
3537+
3538+ #[ tokio:: test]
3539+ async fn test_nullable_unique_build_nonnullable_probe_unbounded_limit ( ) -> Result < ( ) > {
3540+ for join_type in [ JoinType :: Inner , JoinType :: Right , JoinType :: Left ] {
3541+ check_unique_build_join_unbounded_limit (
3542+ join_type,
3543+ false ,
3544+ NullEquality :: NullEqualsNull ,
3545+ )
3546+ . await ?;
3547+ }
3548+ Ok ( ( ) )
35163549}
35173550
35183551fn unique_join_build ( ) -> Result < Arc < DataSourceExec > > {
@@ -3541,6 +3574,7 @@ fn unique_join_build() -> Result<Arc<DataSourceExec>> {
35413574async fn check_unique_build_join_unbounded_limit (
35423575 join_type : JoinType ,
35433576 computed_key : bool ,
3577+ null_equality : NullEquality ,
35443578) -> Result < ( ) > {
35453579 #[ derive( Debug ) ]
35463580 struct PendingPartition ( RecordBatch ) ;
@@ -3558,9 +3592,25 @@ async fn check_unique_build_join_unbounded_limit(
35583592 }
35593593
35603594 let build = unique_join_build ( ) ?;
3595+ let build =
3596+ if null_equality == NullEquality :: NullEqualsNull {
3597+ assert ! ( build. schema( ) . field( 0 ) . is_nullable( ) ) ;
3598+ Arc :: new ( build. as_ref ( ) . clone ( ) . with_constraints (
3599+ Constraints :: new_unverified ( vec ! [ Constraint :: Unique ( vec![ 0 ] ) ] ) ,
3600+ ) )
3601+ } else {
3602+ build
3603+ } ;
35613604
3562- let probe_batch = record_batch ! ( ( "probe_k" , Int32 , [ 1 , 1 ] ) ) ?;
3563- let probe_schema = probe_batch. schema ( ) ;
3605+ let probe_schema = Arc :: new ( Schema :: new ( vec ! [ Field :: new(
3606+ "probe_k" ,
3607+ DataType :: Int32 ,
3608+ false ,
3609+ ) ] ) ) ;
3610+ let probe_batch = RecordBatch :: try_new (
3611+ Arc :: clone ( & probe_schema) ,
3612+ vec ! [ Arc :: new( Int32Array :: from( vec![ 1 , 1 ] ) ) ] ,
3613+ ) ?;
35643614 let probe_ordering = LexOrdering :: new ( [ sort_expr ( "probe_k" , & probe_schema) ] ) . unwrap ( ) ;
35653615 let probe = Arc :: new ( StreamingTableExec :: try_new (
35663616 probe_schema,
@@ -3589,7 +3639,7 @@ async fn check_unique_build_join_unbounded_limit(
35893639 None ,
35903640 join_type,
35913641 vec ! [ SortOptions :: default ( ) ] ,
3592- NullEquality :: NullEqualsNothing ,
3642+ null_equality ,
35933643 ) ?)
35943644 } else {
35953645 Arc :: new (
@@ -3599,6 +3649,7 @@ async fn check_unique_build_join_unbounded_limit(
35993649 vec ! [ ( Arc :: new( Column :: new( "build_k" , 0 ) ) , probe_key) ] ,
36003650 join_type,
36013651 )
3652+ . with_null_equality ( null_equality)
36023653 . with_partition_mode ( PartitionMode :: CollectLeft )
36033654 . build ( ) ?,
36043655 )
@@ -3897,3 +3948,198 @@ async fn test_nullable_unique_build_null_equals_null_requires_sort() -> Result<(
38973948 ) ;
38983949 Ok ( ( ) )
38993950}
3951+
3952+ /// A filter can equate a volatile expression with a column for its own output,
3953+ /// but evaluating the expression again as a join key can produce different values.
3954+ #[ tokio:: test]
3955+ async fn test_unique_build_volatile_key_equivalence_limit ( ) -> Result < ( ) > {
3956+ use datafusion_expr:: { Volatility , create_udf} ;
3957+ use datafusion_physical_expr:: ScalarFunctionExpr ;
3958+ use datafusion_physical_plan:: filter:: FilterExec ;
3959+ use std:: sync:: atomic:: { AtomicUsize , Ordering } ;
3960+
3961+ for original_key_is_volatile in [ true , false ] {
3962+ let calls = AtomicUsize :: new ( 0 ) ;
3963+ let flip = create_udf (
3964+ "volatile_flip" ,
3965+ vec ! [ DataType :: Int32 ] ,
3966+ DataType :: Int32 ,
3967+ Volatility :: Volatile ,
3968+ Arc :: new ( move |args| {
3969+ // The filter keeps every row. On reevaluation in the join, equal
3970+ // probe keys alternate between two different unique build rows.
3971+ if original_key_is_volatile && calls. fetch_add ( 1 , Ordering :: SeqCst ) == 0 {
3972+ return Ok ( args[ 0 ] . clone ( ) ) ;
3973+ }
3974+ let values = if original_key_is_volatile {
3975+ [ 0 , 1 , 0 , 0 ]
3976+ } else {
3977+ [ 1 , 2 , 1 , 2 ]
3978+ } ;
3979+ let len = match & args[ 0 ] {
3980+ ColumnarValue :: Array ( array) => array. len ( ) ,
3981+ ColumnarValue :: Scalar ( _) => 1 ,
3982+ } ;
3983+ Ok ( ColumnarValue :: Array ( Arc :: new (
3984+ Int32Array :: from_iter_values ( values. into_iter ( ) . cycle ( ) . take ( len) ) ,
3985+ ) ) )
3986+ } ) ,
3987+ ) ;
3988+ let probe_batch = record_batch ! (
3989+ ( "probe_k" , Int32 , [ 1 , 1 , 1 , 2 ] ) ,
3990+ ( "probe_q" , Int32 , [ 1 , 2 , 1 , 2 ] )
3991+ ) ?;
3992+ let probe_schema = probe_batch. schema ( ) ;
3993+ let probe_key = col ( "probe_k" , & probe_schema) ?;
3994+ let flip: Arc < dyn PhysicalExpr > = Arc :: new ( ScalarFunctionExpr :: try_new (
3995+ Arc :: new ( flip) ,
3996+ vec ! [ Arc :: clone( & probe_key) ] ,
3997+ & probe_schema,
3998+ Arc :: new ( ConfigOptions :: default ( ) ) ,
3999+ ) ?) ;
4000+ let probe = DataSourceExec :: from_data_source (
4001+ MemorySourceConfig :: try_new (
4002+ & [ vec ! [ probe_batch] ] ,
4003+ Arc :: clone ( & probe_schema) ,
4004+ None ,
4005+ ) ?
4006+ . try_with_sort_information ( vec ! [
4007+ LexOrdering :: new( [ sort_expr( "probe_k" , & probe_schema) ] ) . unwrap( ) ,
4008+ ] ) ?,
4009+ ) ;
4010+ let ( predicate, join_key) : ( _ , Arc < dyn PhysicalExpr > ) =
4011+ if original_key_is_volatile {
4012+ (
4013+ BinaryExpr :: new (
4014+ Arc :: clone ( & probe_key) ,
4015+ Operator :: Eq ,
4016+ Arc :: clone ( & flip) ,
4017+ ) ,
4018+ Arc :: new ( BinaryExpr :: new ( probe_key, Operator :: Plus , flip) ) ,
4019+ )
4020+ } else {
4021+ let q = col ( "probe_q" , & probe_schema) ?;
4022+ ( BinaryExpr :: new ( flip, Operator :: Eq , Arc :: clone ( & q) ) , q)
4023+ } ;
4024+ let probe = Arc :: new ( FilterExec :: try_new ( Arc :: new ( predicate) , probe) ?) ;
4025+ // Equivalence can either hide a volatile child or introduce one through an alias.
4026+ let normalized = probe
4027+ . properties ( )
4028+ . equivalence_properties ( )
4029+ . eq_group ( )
4030+ . normalize_expr ( Arc :: clone ( & join_key) ) ;
4031+ assert_eq ! (
4032+ datafusion_physical_expr_common:: physical_expr:: is_volatile( & normalized) ,
4033+ !original_key_is_volatile,
4034+ ) ;
4035+ let join = HashJoinExecBuilder :: new (
4036+ unique_join_build ( ) ?,
4037+ probe,
4038+ vec ! [ ( Arc :: new( Column :: new( "build_k" , 0 ) ) , join_key) ] ,
4039+ JoinType :: Inner ,
4040+ )
4041+ . with_partition_mode ( PartitionMode :: CollectLeft )
4042+ . build ( ) ?;
4043+ let required = LexOrdering :: new ( [
4044+ sort_expr ( "probe_k" , & join. schema ( ) ) ,
4045+ sort_expr ( "build_k" , & join. schema ( ) ) ,
4046+ sort_expr ( "build_v" , & join. schema ( ) ) ,
4047+ ] )
4048+ . unwrap ( ) ;
4049+ let plan = Arc :: new ( SortExec :: new ( required, Arc :: new ( join) ) . with_fetch ( Some ( 2 ) ) ) ;
4050+ let ctx = SessionContext :: new_with_config (
4051+ SessionConfig :: new ( ) . with_target_partitions ( 1 ) ,
4052+ ) ;
4053+ let batches = datafusion_physical_plan:: collect ( plan, ctx. task_ctx ( ) ) . await ?;
4054+ assert_batches_eq ! (
4055+ [
4056+ "+---------+---------+---------+---------+" ,
4057+ "| build_k | build_v | probe_k | probe_q |" ,
4058+ "+---------+---------+---------+---------+" ,
4059+ "| 1 | 100 | 1 | 1 |" ,
4060+ "| 1 | 100 | 1 | 1 |" ,
4061+ "+---------+---------+---------+---------+" ,
4062+ ] ,
4063+ & batches
4064+ ) ;
4065+ }
4066+
4067+ Ok ( ( ) )
4068+ }
4069+
4070+ #[ tokio:: test]
4071+ async fn test_unique_build_volatile_build_key_equivalence_limit ( ) -> Result < ( ) > {
4072+ use datafusion_expr:: { Volatility , create_udf} ;
4073+ use datafusion_physical_expr:: ScalarFunctionExpr ;
4074+ use datafusion_physical_plan:: filter:: FilterExec ;
4075+ use std:: sync:: atomic:: { AtomicUsize , Ordering } ;
4076+
4077+ let calls = AtomicUsize :: new ( 0 ) ;
4078+ let flip = create_udf (
4079+ "volatile_build_key" ,
4080+ vec ! [ DataType :: Int32 ] ,
4081+ DataType :: Int32 ,
4082+ Volatility :: Volatile ,
4083+ Arc :: new ( move |args| {
4084+ // The filter sees the unique keys [1, 2], but the join sees [1, 1].
4085+ if calls. fetch_add ( 1 , Ordering :: SeqCst ) == 0 {
4086+ Ok ( args[ 0 ] . clone ( ) )
4087+ } else {
4088+ Ok ( ColumnarValue :: Scalar ( 1_i32 . into ( ) ) )
4089+ }
4090+ } ) ,
4091+ ) ;
4092+ let build = unique_join_build ( ) ?;
4093+ let build_key = col ( "build_k" , & build. schema ( ) ) ?;
4094+ let flip: Arc < dyn PhysicalExpr > = Arc :: new ( ScalarFunctionExpr :: try_new (
4095+ Arc :: new ( flip) ,
4096+ vec ! [ Arc :: clone( & build_key) ] ,
4097+ & build. schema ( ) ,
4098+ Arc :: new ( ConfigOptions :: default ( ) ) ,
4099+ ) ?) ;
4100+ let build = Arc :: new ( FilterExec :: try_new (
4101+ Arc :: new ( BinaryExpr :: new ( build_key, Operator :: Eq , Arc :: clone ( & flip) ) ) ,
4102+ build,
4103+ ) ?) ;
4104+ let probe_batch = record_batch ! ( ( "probe_k" , Int32 , [ 1 , 1 ] ) ) ?;
4105+ let probe_schema = probe_batch. schema ( ) ;
4106+ let probe = DataSourceExec :: from_data_source (
4107+ MemorySourceConfig :: try_new (
4108+ & [ vec ! [ probe_batch] ] ,
4109+ Arc :: clone ( & probe_schema) ,
4110+ None ,
4111+ ) ?
4112+ . try_with_sort_information ( vec ! [
4113+ LexOrdering :: new( [ sort_expr( "probe_k" , & probe_schema) ] ) . unwrap( ) ,
4114+ ] ) ?,
4115+ ) ;
4116+ let join = HashJoinExecBuilder :: new (
4117+ build,
4118+ probe,
4119+ vec ! [ ( flip, col( "probe_k" , & probe_schema) ?) ] ,
4120+ JoinType :: Inner ,
4121+ )
4122+ . with_partition_mode ( PartitionMode :: CollectLeft )
4123+ . build ( ) ?;
4124+ let required = LexOrdering :: new ( [
4125+ sort_expr ( "probe_k" , & join. schema ( ) ) ,
4126+ sort_expr ( "build_k" , & join. schema ( ) ) ,
4127+ sort_expr ( "build_v" , & join. schema ( ) ) ,
4128+ ] )
4129+ . unwrap ( ) ;
4130+ let plan = Arc :: new ( SortExec :: new ( required, Arc :: new ( join) ) . with_fetch ( Some ( 2 ) ) ) ;
4131+ let batches =
4132+ datafusion_physical_plan:: collect ( plan, SessionContext :: new ( ) . task_ctx ( ) ) . await ?;
4133+ assert_batches_eq ! (
4134+ [
4135+ "+---------+---------+---------+" ,
4136+ "| build_k | build_v | probe_k |" ,
4137+ "+---------+---------+---------+" ,
4138+ "| 1 | 100 | 1 |" ,
4139+ "| 1 | 100 | 1 |" ,
4140+ "+---------+---------+---------+" ,
4141+ ] ,
4142+ & batches
4143+ ) ;
4144+ Ok ( ( ) )
4145+ }
0 commit comments