Skip to content

Commit 4abe3d0

Browse files
authored
fix: TableStrategy currently hides panic in spawned column writer tasks. (#8672)
## Rationale for this change `TableStrategy` currently hides panics that happen in the individual column streams. ## What changes are included in this PR? Instead of detaching the fanout work, we await it and make sure everything finishes successfully. ## What APIs are changed? Are there any user-facing changes? No API change, added test to verify panic propagation. Signed-off-by: Adam Gutglick <adam@spiraldb.com>
1 parent 9a42615 commit 4abe3d0

1 file changed

Lines changed: 76 additions & 19 deletions

File tree

vortex-layout/src/layouts/table.rs

Lines changed: 76 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ use std::sync::Arc;
1010
use async_trait::async_trait;
1111
use futures::StreamExt;
1212
use futures::TryStreamExt;
13+
use futures::future::try_join;
1314
use futures::future::try_join_all;
1415
use futures::pin_mut;
1516
use 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)]
371375
mod 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

Comments
 (0)