@@ -398,6 +398,127 @@ async fn roundtrip_table_scan_offset_and_fetch() -> Result<()> {
398398 Ok ( ( ) )
399399}
400400
401+ // Limit expressions use Int64, but limit pushdown can produce a larger usize fetch.
402+ #[ tokio:: test]
403+ #[ cfg( target_pointer_width = "64" ) ]
404+ async fn roundtrip_table_scan_fetch_above_i64_max ( ) -> Result < ( ) > {
405+ use datafusion:: arrow:: array:: Int64Array ;
406+ use datafusion:: arrow:: record_batch:: RecordBatch ;
407+ use datafusion:: dataframe:: DataFrame ;
408+
409+ let ctx = SessionContext :: new ( ) ;
410+ ctx. register_batch (
411+ "t" ,
412+ RecordBatch :: try_from_iter ( vec ! [ (
413+ "x" ,
414+ Arc :: new( Int64Array :: from( vec![ 0 , 1 , 2 ] ) ) as ArrayRef ,
415+ ) ] ) ?,
416+ ) ?;
417+ let plan = ctx
418+ . sql ( "SELECT * FROM t LIMIT 9223372036854775807 OFFSET 1" )
419+ . await ?
420+ . into_optimized_plan ( ) ?;
421+ let LogicalPlan :: Limit ( limit) = & plan else {
422+ panic ! ( "expected Limit" )
423+ } ;
424+ let LogicalPlan :: TableScan ( scan) = limit. input . as_ref ( ) else {
425+ panic ! ( "expected TableScan" )
426+ } ;
427+ assert_eq ! ( scan. fetch, Some ( i64 :: MAX as usize + 1 ) ) ;
428+ let proto = to_substrait_plan ( & plan, & ctx. state ( ) ) ?;
429+ let restored = from_substrait_plan ( & ctx. state ( ) , & proto) . await ?;
430+ for plan in [ plan, restored] {
431+ let batches = DataFrame :: new ( ctx. state ( ) , plan) . collect ( ) . await ?;
432+ datafusion:: assert_batches_eq!(
433+ [ "+---+" , "| x |" , "+---+" , "| 1 |" , "| 2 |" , "+---+" ] ,
434+ & batches
435+ ) ;
436+ }
437+ Ok ( ( ) )
438+ }
439+
440+ #[ tokio:: test]
441+ #[ cfg( target_pointer_width = "64" ) ]
442+ async fn roundtrip_table_scan_signed_range ( ) -> Result < ( ) > {
443+ use datafusion:: arrow:: array:: Int64Array ;
444+ use datafusion:: dataframe:: DataFrame ;
445+ use datafusion:: datasource:: provider_as_source;
446+ use datafusion:: logical_expr:: TableScanBuilder ;
447+
448+ let ctx = SessionContext :: new ( ) ;
449+ let temp = tempfile:: tempdir ( ) ?;
450+ let path = temp. path ( ) . join ( "data.csv" ) ;
451+ std:: fs:: write ( & path, "x\n 0\n 1\n 2\n " ) ?;
452+ ctx. register_csv ( "t" , path. to_str ( ) . unwrap ( ) , CsvReadOptions :: new ( ) )
453+ . await ?;
454+ let source = provider_as_source ( ctx. table_provider ( "t" ) . await ?) ;
455+ let max = i64:: MAX as usize ;
456+ for ( skip, fetch, parent, succeeds) in [
457+ ( None , Some ( max) , None , true ) ,
458+ ( Some ( max) , None , None , true ) ,
459+ ( None , Some ( max + 1 ) , None , false ) ,
460+ ( None , Some ( usize:: MAX ) , None , false ) ,
461+ ( Some ( max + 1 ) , None , None , false ) ,
462+ ( Some ( usize:: MAX ) , None , None , false ) ,
463+ // The parent covers the scan bound exactly, or imposes a tighter bound.
464+ ( None , Some ( max + 1 ) , Some ( ( 1 , Some ( max) ) ) , true ) ,
465+ ( None , Some ( usize:: MAX ) , Some ( ( 1 , Some ( 2 ) ) ) , true ) ,
466+ ( Some ( 1 ) , Some ( max + 1 ) , Some ( ( 0 , Some ( 2 ) ) ) , true ) ,
467+ // These parents do not make the scan fetch redundant.
468+ ( None , Some ( max + 1 ) , Some ( ( max, Some ( max) ) ) , false ) ,
469+ ( None , Some ( max + 1 ) , Some ( ( 1 , None ) ) , false ) ,
470+ ] {
471+ let scan = TableScanBuilder :: new ( "t" , source. clone ( ) )
472+ . with_skip ( skip)
473+ . with_fetch ( fetch)
474+ . build ( ) ?;
475+ let mut builder = LogicalPlanBuilder :: from ( LogicalPlan :: TableScan ( scan) ) ;
476+ if let Some ( ( skip, fetch) ) = parent {
477+ builder = builder. limit ( skip, fetch) ?;
478+ }
479+ let plan = builder. build ( ) ?;
480+ let proto = to_substrait_plan ( & plan, & ctx. state ( ) ) ;
481+ if !succeeds {
482+ let error = proto. unwrap_err ( ) . to_string ( ) ;
483+ let field = if skip. is_some ( ) { "offset" } else { "fetch" } ;
484+ assert ! (
485+ error. contains( & format!(
486+ "TableScan {field} exceeds the supported Int64 range"
487+ ) ) ,
488+ "{error}"
489+ ) ;
490+ continue ;
491+ }
492+ let proto = proto?;
493+ let restored = from_substrait_plan ( & ctx. state ( ) , & proto) . await ?;
494+ let expected = if skip == Some ( max) {
495+ vec ! [ ]
496+ } else if parent. is_some ( ) {
497+ vec ! [ 1 , 2 ]
498+ } else {
499+ vec ! [ 0 , 1 , 2 ]
500+ } ;
501+ for plan in [ plan, restored] {
502+ let batches = DataFrame :: new ( ctx. state ( ) , plan) . collect ( ) . await ?;
503+ let actual = batches
504+ . iter ( )
505+ . flat_map ( |batch| {
506+ batch
507+ . column ( 0 )
508+ . as_any ( )
509+ . downcast_ref :: < Int64Array > ( )
510+ . unwrap ( )
511+ . values ( )
512+ . iter ( )
513+ . copied ( )
514+ } )
515+ . collect :: < Vec < _ > > ( ) ;
516+ assert_eq ! ( actual, expected) ;
517+ }
518+ }
519+ Ok ( ( ) )
520+ }
521+
401522#[ tokio:: test]
402523async fn select_with_limit_offset ( ) -> Result < ( ) > {
403524 roundtrip ( "SELECT * FROM data LIMIT 200 OFFSET 10" ) . await ?;
0 commit comments