@@ -268,6 +268,12 @@ impl ArrayParentReduceRule<Dict> for DictionaryScalarFnCodesPullUpRule {
268268 return Ok ( None ) ;
269269 }
270270
271+ // Pulling up the codes evaluates every dictionary value, including unreferenced entries
272+ // on which a fallible function could fail even though no row would evaluate them.
273+ if !array. has_all_values_referenced ( ) && !parent. scalar_fn ( ) . signature ( ) . is_infallible ( ) {
274+ return Ok ( None ) ;
275+ }
276+
271277 // Check that all siblings are dictionaries, and have the same number of values as us.
272278 // This is a cheap first loop.
273279 if !parent. iter_children ( ) . enumerate ( ) . all ( |( idx, c) | {
@@ -310,11 +316,13 @@ impl ArrayParentReduceRule<Dict> for DictionaryScalarFnCodesPullUpRule {
310316mod tests {
311317 use std:: sync:: LazyLock ;
312318
319+ use rstest:: rstest;
313320 use vortex_buffer:: buffer;
314321 use vortex_error:: VortexResult ;
315322 use vortex_session:: VortexSession ;
316323
317324 use crate :: ArrayRef ;
325+ use crate :: Canonical ;
318326 use crate :: IntoArray ;
319327 use crate :: arrays:: BoolArray ;
320328 use crate :: arrays:: Chunked ;
@@ -444,4 +452,77 @@ mod tests {
444452 assert_arrays_eq ! ( optimized, expected, & mut ctx) ;
445453 Ok ( ( ) )
446454 }
455+
456+ /// Both slices share the codes `[1, 1, 1]`, so the codes pull-up divides the two values
457+ /// arrays. The divisor still holds the unreferenced `0`, which no row divides by, so the
458+ /// result must be `[8, 8, 8] / [2, 2, 2]`.
459+ #[ test]
460+ fn codes_pull_up_skips_fallible_fn_over_unreferenced_values ( ) -> VortexResult < ( ) > {
461+ let mut ctx = SESSION . create_execution_ctx ( ) ;
462+ let dividend = DictArray :: try_new (
463+ buffer ! [ 0u8 , 1 , 1 , 1 ] . into_array ( ) ,
464+ buffer ! [ 7i32 , 8 ] . into_array ( ) ,
465+ ) ?
466+ . into_array ( )
467+ . slice ( 1 ..4 ) ?
468+ . optimize_ctx ( & SESSION ) ?;
469+ let divisor = DictArray :: try_new (
470+ buffer ! [ 0u8 , 1 , 1 , 1 ] . into_array ( ) ,
471+ buffer ! [ 0i32 , 2 ] . into_array ( ) ,
472+ ) ?
473+ . into_array ( )
474+ . slice ( 1 ..4 ) ?
475+ . optimize_ctx ( & SESSION ) ?;
476+ let expr = Binary :: try_new ( dividend, divisor, Operator :: Div ) ?. into_array ( ) ;
477+
478+ let quotient = expr
479+ . optimize_ctx ( & SESSION ) ?
480+ . execute :: < Canonical > ( & mut ctx) ?;
481+
482+ assert_arrays_eq ! ( quotient, PrimitiveArray :: from_iter( [ 4i32 , 4 , 4 ] ) , & mut ctx) ;
483+ Ok ( ( ) )
484+ }
485+
486+ #[ rstest]
487+ #[ case:: infallible( false , Operator :: Eq , BoolArray :: from_iter( [ false , true , false , true ] ) . into_array( ) ) ]
488+ #[ case:: all_values_referenced( true , Operator :: Div , buffer![ 4i32 , 1 , 4 , 1 ] . into_array( ) ) ]
489+ fn codes_pull_up_preserves_safe_rewrites (
490+ #[ case] all_values_referenced : bool ,
491+ #[ case] operator : Operator ,
492+ #[ case] expected : ArrayRef ,
493+ ) -> VortexResult < ( ) > {
494+ let codes = buffer ! [ 0u8 , 1 , 0 , 1 ] . into_array ( ) ;
495+ // SAFETY: Both dictionary entries are referenced by valid codes.
496+ let lhs = unsafe {
497+ DictArray :: try_new ( codes. clone ( ) , buffer ! [ 8i32 , 9 ] . into_array ( ) ) ?
498+ . set_all_values_referenced ( all_values_referenced)
499+ }
500+ . into_array ( ) ;
501+ let rhs = DictArray :: try_new ( codes, buffer ! [ 2i32 , 9 ] . into_array ( ) ) ?. into_array ( ) ;
502+
503+ let optimized = Binary :: try_new ( lhs, rhs, operator) ?
504+ . into_array ( )
505+ . optimize_ctx ( & SESSION ) ?;
506+ assert ! ( optimized. is:: <Dict >( ) ) ;
507+ let mut ctx = SESSION . create_execution_ctx ( ) ;
508+ let result = optimized. execute :: < Canonical > ( & mut ctx) ?;
509+ assert_arrays_eq ! ( result, expected, & mut ctx) ;
510+ Ok ( ( ) )
511+ }
512+
513+ #[ test]
514+ fn codes_pull_up_rejects_referenced_zero_divisor ( ) -> VortexResult < ( ) > {
515+ let codes = buffer ! [ 0u8 , 1 , 0 , 1 ] . into_array ( ) ;
516+ let dividend =
517+ DictArray :: try_new ( codes. clone ( ) , buffer ! [ 7i32 , 8 ] . into_array ( ) ) ?. into_array ( ) ;
518+ let divisor = DictArray :: try_new ( codes, buffer ! [ 0i32 , 2 ] . into_array ( ) ) ?. into_array ( ) ;
519+ let expr = Binary :: try_new ( dividend, divisor, Operator :: Div ) ?. into_array ( ) ;
520+
521+ let error = expr
522+ . optimize_ctx ( & SESSION ) ?
523+ . execute :: < Canonical > ( & mut SESSION . create_execution_ctx ( ) )
524+ . unwrap_err ( ) ;
525+ assert ! ( error. to_string( ) . contains( "division by zero" ) ) ;
526+ Ok ( ( ) )
527+ }
447528}
0 commit comments