@@ -320,6 +320,84 @@ async fn select_without_limit() -> Result<()> {
320320 roundtrip_fill_na ( "SELECT * FROM data OFFSET 5+7-2" ) . await
321321}
322322
323+ #[ tokio:: test]
324+ async fn roundtrip_table_scan_offset_and_fetch ( ) -> Result < ( ) > {
325+ use datafusion:: arrow:: array:: Int64Array ;
326+ use datafusion:: dataframe:: DataFrame ;
327+ use datafusion:: datasource:: provider_as_source;
328+ use datafusion:: logical_expr:: TableScanBuilder ;
329+
330+ // A single CSV partition preserves input order and supports exact skip pushdown.
331+ let temp = tempfile:: tempdir ( ) ?;
332+ let path = temp. path ( ) . join ( "data.csv" ) ;
333+ std:: fs:: write ( & path, "x\n 0\n 1\n 2\n 3\n 4\n 5\n " ) ?;
334+ let ctx = SessionContext :: new ( ) ;
335+ ctx. register_csv ( "t" , path. to_str ( ) . unwrap ( ) , CsvReadOptions :: new ( ) )
336+ . await ?;
337+ let source = provider_as_source ( ctx. table_provider ( "t" ) . await ?) ;
338+
339+ for skip in [ Some ( 2 ) , None , Some ( 0 ) , Some ( 8 ) ] {
340+ for fetch in [ None , Some ( 0 ) , Some ( 2 ) ] {
341+ for project in [ false , true ] {
342+ for parent in [
343+ None ,
344+ Some ( ( 0 , Some ( 2 ) ) ) ,
345+ Some ( ( 1 , Some ( 3 ) ) ) ,
346+ Some ( ( 1 , None ) ) ,
347+ ] {
348+ let scan = TableScanBuilder :: new ( "t" , source. clone ( ) )
349+ . with_skip ( skip)
350+ . with_fetch ( fetch)
351+ . build ( ) ?;
352+ let mut builder =
353+ LogicalPlanBuilder :: from ( LogicalPlan :: TableScan ( scan) ) ;
354+ if project {
355+ builder = builder. project ( vec ! [ col( "x" ) ] ) ?;
356+ }
357+ if let Some ( ( skip, fetch) ) = parent {
358+ builder = builder. limit ( skip, fetch) ?;
359+ }
360+ let plan = builder. build ( ) ?;
361+ let proto = to_substrait_plan ( & plan, & ctx. state ( ) ) ?;
362+ let restored = from_substrait_plan ( & ctx. state ( ) , & proto) . await ?;
363+ let mut expected = ( 0_i64 ..6 )
364+ . skip ( skip. unwrap_or ( 0 ) )
365+ . take ( fetch. unwrap_or ( usize:: MAX ) )
366+ . collect :: < Vec < _ > > ( ) ;
367+ if let Some ( ( skip, fetch) ) = parent {
368+ expected = expected
369+ . into_iter ( )
370+ . skip ( skip)
371+ . take ( fetch. unwrap_or ( usize:: MAX ) )
372+ . collect ( ) ;
373+ }
374+ for ( stage, plan) in [ ( "original" , plan) , ( "roundtrip" , restored) ] {
375+ let batches = DataFrame :: new ( ctx. state ( ) , plan) . collect ( ) . await ?;
376+ let actual = batches
377+ . iter ( )
378+ . flat_map ( |batch| {
379+ batch
380+ . column ( 0 )
381+ . as_any ( )
382+ . downcast_ref :: < Int64Array > ( )
383+ . unwrap ( )
384+ . values ( )
385+ . iter ( )
386+ . copied ( )
387+ } )
388+ . collect :: < Vec < _ > > ( ) ;
389+ assert_eq ! (
390+ actual, expected,
391+ "{stage}: skip={skip:?}, fetch={fetch:?}, project={project}, parent={parent:?}"
392+ ) ;
393+ }
394+ }
395+ }
396+ }
397+ }
398+ Ok ( ( ) )
399+ }
400+
323401#[ tokio:: test]
324402async fn select_with_limit_offset ( ) -> Result < ( ) > {
325403 roundtrip ( "SELECT * FROM data LIMIT 200 OFFSET 10" ) . await ?;
0 commit comments