@@ -41,6 +41,7 @@ use vortex::file::v2::FileStatsLayoutReader;
4141use vortex:: io:: kanal_ext:: KanalExt as _;
4242use vortex:: io:: runtime:: BlockingRuntime as _;
4343use vortex:: io:: runtime:: current:: ThreadSafeIterator ;
44+ use vortex:: io:: runtime:: current:: TryRecv ;
4445use vortex:: layout:: scan:: multi:: MultiLayoutChild ;
4546use vortex:: layout:: scan:: multi:: MultiLayoutDataSource ;
4647use vortex:: metrics:: tracing:: get_global_labels;
@@ -131,6 +132,7 @@ impl<'a> TableInitInput<'a> {
131132
132133type ScanItem = VortexResult < ( ArrayRef , Arc < ConversionCache > ) > ;
133134type DataSourceIterator = ThreadSafeIterator < ScanItem > ;
135+ type ScanItemFuture = BoxFuture < ' static , Option < ScanItem > > ;
134136
135137pub struct TableFunctionGlobal {
136138 iterator : DataSourceIterator ,
@@ -158,8 +160,10 @@ pub struct TableFunctionLocal {
158160 exporter : Option < ArrayExporter > ,
159161 partition_index : u64 ,
160162 file_index : usize ,
161- // Aggregate scan accumulated partials. Empty for non-aggregate scan
163+ /// Aggregate scan accumulated partials. Empty for non-aggregate scan
162164 partials : Vec < Box < dyn DynAccumulator > > ,
165+ /// In progress next item
166+ pending : Option < ScanItemFuture > ,
163167}
164168
165169pub struct PartitionData {
@@ -422,6 +426,7 @@ pub fn init_local(
422426 partition_index : 0 ,
423427 file_index : 0 ,
424428 partials,
429+ pending : None ,
425430 }
426431}
427432
@@ -443,11 +448,43 @@ fn convert_result(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Struc
443448 } )
444449}
445450
451+ /// Write aggregation output row to "chunk"
452+ fn finalize_aggregate (
453+ global_state : & TableFunctionGlobal ,
454+ chunk : & mut DataChunkRef ,
455+ ) -> VortexResult < ( ) > {
456+ // 0 means we're the last thread, u64::MAX means output is written.
457+ // is_err() means CAS didn't succeed
458+ if global_state
459+ . pending
460+ . compare_exchange ( 0 , u64:: MAX , Ordering :: AcqRel , Ordering :: Relaxed )
461+ . is_err ( )
462+ {
463+ return Ok ( ( ) ) ;
464+ }
465+
466+ let mut accumulators = global_state. partials . lock ( ) ;
467+ let row_count = global_state. row_count . load ( Ordering :: Acquire ) as i64 ;
468+ let mut accum_iter = accumulators. iter_mut ( ) ;
469+ for ( idx, aggregate) in global_state. aggregates . iter ( ) . enumerate ( ) {
470+ let value = match aggregate {
471+ ColumnAggregate :: Real { .. } => {
472+ let accum = accum_iter. next ( ) . vortex_expect ( "partial for real agg" ) ;
473+ Value :: try_from ( accum. finish ( ) ?) ?
474+ }
475+ ColumnAggregate :: CountStar => Value :: from ( row_count) ,
476+ } ;
477+ chunk. get_vector_mut ( idx) . reference_value ( & value) ;
478+ }
479+ chunk. set_len ( 1 ) ;
480+ Ok ( ( ) )
481+ }
482+
446483fn scan_aggregate (
447484 local_state : & mut TableFunctionLocal ,
448485 global_state : & TableFunctionGlobal ,
449486 chunk : & mut DataChunkRef ,
450- ) -> VortexResult < ( ) > {
487+ ) -> VortexResult < ScanResult > {
451488 let aggregates_len = global_state. aggregates . len ( ) ;
452489 // seen[k] = output column for requested column k.
453490 // If min(x), max(x), avg(y) are requested, seen = { 0: 0, 1: 1}
@@ -468,32 +505,13 @@ fn scan_aggregate(
468505
469506 let mut ctx = SESSION . create_execution_ctx ( ) ;
470507 loop {
471- let Some ( result) = local_state. iterator . next ( ) else {
472- // 0 means we're the last thread, u64::MAX means output is written.
473- // is_err() means CAS didn't succeed
474- if global_state
475- . pending
476- . compare_exchange ( 0 , u64:: MAX , Ordering :: AcqRel , Ordering :: Relaxed )
477- . is_err ( )
478- {
479- return Ok ( ( ) ) ;
508+ let result = match poll_next_item ( local_state) {
509+ NextItem :: Ready ( item) => item,
510+ NextItem :: BlockedOnIO => return Ok ( ScanResult :: BlockedOnIO ) ,
511+ NextItem :: None => {
512+ finalize_aggregate ( global_state, chunk) ?;
513+ return Ok ( ScanResult :: Exported ) ;
480514 }
481-
482- let mut accumulators = global_state. partials . lock ( ) ;
483- let row_count = global_state. row_count . load ( Ordering :: Acquire ) as i64 ;
484- let mut accum_iter = accumulators. iter_mut ( ) ;
485- for ( idx, aggregate) in global_state. aggregates . iter ( ) . enumerate ( ) {
486- let value = match aggregate {
487- ColumnAggregate :: Real { .. } => {
488- let accum = accum_iter. next ( ) . vortex_expect ( "partial for real agg" ) ;
489- Value :: try_from ( accum. finish ( ) ?) ?
490- }
491- ColumnAggregate :: CountStar => Value :: from ( row_count) ,
492- } ;
493- chunk. get_vector_mut ( idx) . reference_value ( & value) ;
494- }
495- chunk. set_len ( 1 ) ;
496- return Ok ( ( ) ) ;
497515 } ;
498516 let array = convert_result ( result?. 0 , & mut ctx) ?;
499517
@@ -517,22 +535,77 @@ fn scan_aggregate(
517535 }
518536}
519537
538+ enum NextItem {
539+ Ready ( ScanItem ) ,
540+ None ,
541+ BlockedOnIO ,
542+ }
543+
544+ /// Poll next item without blocking caller thread
545+ fn poll_next_item ( local_state : & mut TableFunctionLocal ) -> NextItem {
546+ if let Some ( future) = local_state. pending . as_mut ( ) {
547+ let mut ctx = Context :: from_waker ( std:: task:: Waker :: noop ( ) ) ;
548+ return match future. poll_unpin ( & mut ctx) {
549+ Poll :: Ready ( item) => {
550+ local_state. pending = None ;
551+ if let Some ( item) = item {
552+ NextItem :: Ready ( item)
553+ } else {
554+ NextItem :: None
555+ }
556+ }
557+ Poll :: Pending => NextItem :: BlockedOnIO ,
558+ } ;
559+ }
560+
561+ match local_state. iterator . try_recv ( ) {
562+ TryRecv :: Item ( item) => NextItem :: Ready ( item) ,
563+ TryRecv :: Closed => NextItem :: None ,
564+ TryRecv :: Empty => {
565+ let rx = local_state. iterator . receiver ( ) ;
566+ local_state. pending = Some ( Box :: pin ( async move { rx. recv ( ) . await . ok ( ) } ) ) ;
567+ NextItem :: BlockedOnIO
568+ }
569+ }
570+ }
571+
572+ /// Finish array receive started by scan()
573+ pub fn wait_and_fetch ( local_state : & mut TableFunctionLocal ) {
574+ if let Some ( mut future) = local_state. pending . take ( ) {
575+ // Duckdb calls this on a background task thread pool which is separate
576+ // from table function processing pool. This call therefore doesn't
577+ // block the query worker
578+ let item = RUNTIME . block_on ( & mut future) ;
579+ // When this assignment is happening, thread related to "local_state"
580+ // is suspended. It will be resumed only when wait_and_fetch returns
581+ // so there's no race condition
582+ local_state. pending = Some ( Box :: pin ( std:: future:: ready ( item) ) ) ;
583+ } ;
584+ }
585+
586+ #[ derive( PartialEq , Eq ) ]
587+ pub enum ScanResult {
588+ Exported ,
589+ BlockedOnIO ,
590+ }
591+
520592pub fn scan (
521593 local_state : & mut TableFunctionLocal ,
522594 global_state : & TableFunctionGlobal ,
523595 chunk : & mut DataChunkRef ,
524- ) -> VortexResult < ( ) > {
596+ ) -> VortexResult < ScanResult > {
525597 if !local_state. partials . is_empty ( ) {
526598 return scan_aggregate ( local_state, global_state, chunk) ;
527599 }
528600
529601 loop {
530602 if local_state. exporter . is_none ( ) {
531603 let mut ctx = SESSION . create_execution_ctx ( ) ;
532- let Some ( result) = local_state. iterator . next ( ) else {
533- return Ok ( ( ) ) ;
604+ let ( array_result, conversion_cache) = match poll_next_item ( local_state) {
605+ NextItem :: Ready ( item) => item?,
606+ NextItem :: None => return Ok ( ScanResult :: Exported ) ,
607+ NextItem :: BlockedOnIO => return Ok ( ScanResult :: BlockedOnIO ) ,
534608 } ;
535- let ( array_result, conversion_cache) = result?;
536609 local_state. file_index = conversion_cache. file_index ;
537610 let array_result = convert_result ( array_result, & mut ctx) ?;
538611
@@ -576,7 +649,7 @@ pub fn scan(
576649 . reference_value ( & Value :: from ( local_state. file_index as u64 ) ) ;
577650 }
578651
579- Ok ( ( ) )
652+ Ok ( ScanResult :: Exported )
580653}
581654
582655/// Scan progress as a percentage (0.0–100.0).
0 commit comments