@@ -10,6 +10,7 @@ use std::sync::Arc;
1010use async_trait:: async_trait;
1111use futures:: StreamExt ;
1212use futures:: TryStreamExt ;
13+ use futures:: future:: try_join;
1314use futures:: future:: try_join_all;
1415use futures:: pin_mut;
1516use itertools:: Itertools ;
@@ -266,30 +267,33 @@ impl LayoutStrategy for TableStrategy {
266267 let ( column_streams_tx, column_streams_rx) : ( Vec < _ > , Vec < _ > ) =
267268 ( 0 ..stream_count) . map ( |_| kanal:: bounded_async ( 1 ) ) . unzip ( ) ;
268269
269- // Spawn a task to fan out column chunks to their respective transposed streams
270+ // Fan out column chunks to their respective transposed streams. Keep this future joined
271+ // with the column writers so producer panics/errors cannot be hidden as channel EOF.
270272 let handle = session. handle ( ) ;
271- handle
272- . spawn ( async move {
273- pin_mut ! ( columns_vec_stream) ;
274- while let Some ( result) = columns_vec_stream. next ( ) . await {
275- match result {
276- Ok ( columns) => {
277- for ( tx, column) in column_streams_tx. iter ( ) . zip_eq ( columns. into_iter ( ) )
278- {
279- let _ = tx. send ( Ok ( column) ) . await ;
273+ let fanout_fut = async move {
274+ pin_mut ! ( columns_vec_stream) ;
275+ while let Some ( result) = columns_vec_stream. next ( ) . await {
276+ match result {
277+ Ok ( columns) => {
278+ for ( tx, column) in column_streams_tx. iter ( ) . zip_eq ( columns. into_iter ( ) ) {
279+ if tx. send ( Ok ( column) ) . await . is_err ( ) {
280+ vortex_bail ! (
281+ "table column writer finished before all chunks were sent"
282+ ) ;
280283 }
281284 }
282- Err ( e) => {
283- let e: Arc < VortexError > = Arc :: new ( e) ;
284- for tx in column_streams_tx. iter ( ) {
285- let _ = tx. send ( Err ( VortexError :: from ( Arc :: clone ( & e) ) ) ) . await ;
286- }
287- break ;
285+ }
286+ Err ( e) => {
287+ let e: Arc < VortexError > = Arc :: new ( e) ;
288+ for tx in column_streams_tx. iter ( ) {
289+ let _ = tx. send ( Err ( VortexError :: from ( Arc :: clone ( & e) ) ) ) . await ;
288290 }
291+ return Err ( VortexError :: from ( e) ) ;
289292 }
290293 }
291- } )
292- . detach ( ) ;
294+ }
295+ Ok ( ( ) )
296+ } ;
293297
294298 // First child column is the validity, subsequence children are the individual struct fields
295299 let column_dtypes: Vec < DType > = if is_nullable {
@@ -359,7 +363,7 @@ impl LayoutStrategy for TableStrategy {
359363 } )
360364 . collect ( ) ;
361365
362- let column_layouts = try_join_all ( layout_futures) . await ?;
366+ let ( _success , column_layouts) = try_join ( fanout_fut , try_join_all ( layout_futures) ) . await ?;
363367 // TODO(os): transposed stream could count row counts as well,
364368 // This must hold though, all columns must have the same row count of the struct layout
365369 let row_count = column_layouts. first ( ) . map ( |l| l. row_count ( ) ) . unwrap_or ( 0 ) ;
@@ -370,12 +374,28 @@ impl LayoutStrategy for TableStrategy {
370374#[ cfg( test) ]
371375mod tests {
372376 use std:: sync:: Arc ;
377+ use std:: task:: Poll ;
373378
379+ use vortex_array:: ArrayContext ;
380+ use vortex_array:: ArrayRef ;
381+ use vortex_array:: dtype:: DType ;
374382 use vortex_array:: dtype:: FieldPath ;
383+ use vortex_array:: dtype:: Nullability ;
384+ use vortex_array:: dtype:: PType ;
385+ use vortex_array:: dtype:: StructFields ;
375386 use vortex_array:: field_path;
387+ use vortex_error:: VortexResult ;
388+ use vortex_io:: runtime:: single:: block_on;
389+ use vortex_io:: session:: RuntimeSessionExt ;
376390
391+ use crate :: LayoutStrategy ;
377392 use crate :: layouts:: flat:: writer:: FlatLayoutStrategy ;
378393 use crate :: layouts:: table:: TableStrategy ;
394+ use crate :: segments:: TestSegments ;
395+ use crate :: sequence:: SequenceId ;
396+ use crate :: sequence:: SequentialStreamAdapter ;
397+ use crate :: sequence:: SequentialStreamExt ;
398+ use crate :: test:: SESSION ;
379399
380400 #[ test]
381401 #[ should_panic(
@@ -407,4 +427,41 @@ mod tests {
407427 )
408428 . with_field_writer ( FieldPath :: root ( ) , flat) ;
409429 }
430+
431+ #[ test]
432+ #[ should_panic( expected = "panic while transposing table stream" ) ]
433+ fn table_fanout_panic_propagates ( ) {
434+ let ctx = ArrayContext :: empty ( ) ;
435+ let segments = Arc :: new ( TestSegments :: default ( ) ) ;
436+ let ( _, eof) = SequenceId :: root ( ) . split ( ) ;
437+ let dtype = DType :: Struct (
438+ StructFields :: from_iter ( [ (
439+ "a" ,
440+ DType :: Primitive ( PType :: I32 , Nullability :: NonNullable ) ,
441+ ) ] ) ,
442+ Nullability :: NonNullable ,
443+ ) ;
444+ let stream =
445+ futures:: stream:: poll_fn ( |_| -> Poll < Option < VortexResult < ( SequenceId , ArrayRef ) > > > {
446+ panic ! ( "panic while transposing table stream" ) ;
447+ } ) ;
448+ let strategy = TableStrategy :: new (
449+ Arc :: new ( FlatLayoutStrategy :: default ( ) ) ,
450+ Arc :: new ( FlatLayoutStrategy :: default ( ) ) ,
451+ ) ;
452+
453+ block_on ( |handle| async move {
454+ let session = SESSION . clone ( ) . with_handle ( handle) ;
455+ strategy
456+ . write_stream (
457+ ctx,
458+ segments,
459+ SequentialStreamAdapter :: new ( dtype, stream) . sendable ( ) ,
460+ eof,
461+ & session,
462+ )
463+ . await
464+ . unwrap ( ) ;
465+ } ) ;
466+ }
410467}
0 commit comments