From 68badd6899c82156a3a6d4c4cf9906051a674ece Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 7 Jul 2026 17:29:51 +0100 Subject: [PATCH] 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 (cherry picked from commit 4abe3d04f097b1fbbfcae23795f347b7af5b0919) --- vortex-layout/src/layouts/table.rs | 95 ++++++++++++++++++++++++------ 1 file changed, 76 insertions(+), 19 deletions(-) diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 28b3af52c87..04c47a9f732 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -10,6 +10,7 @@ use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt; use futures::TryStreamExt; +use futures::future::try_join; use futures::future::try_join_all; use futures::pin_mut; use itertools::Itertools; @@ -266,30 +267,33 @@ impl LayoutStrategy for TableStrategy { let (column_streams_tx, column_streams_rx): (Vec<_>, Vec<_>) = (0..stream_count).map(|_| kanal::bounded_async(1)).unzip(); - // Spawn a task to fan out column chunks to their respective transposed streams + // Fan out column chunks to their respective transposed streams. Keep this future joined + // with the column writers so producer panics/errors cannot be hidden as channel EOF. let handle = session.handle(); - handle - .spawn(async move { - pin_mut!(columns_vec_stream); - while let Some(result) = columns_vec_stream.next().await { - match result { - Ok(columns) => { - for (tx, column) in column_streams_tx.iter().zip_eq(columns.into_iter()) - { - let _ = tx.send(Ok(column)).await; + let fanout_fut = async move { + pin_mut!(columns_vec_stream); + while let Some(result) = columns_vec_stream.next().await { + match result { + Ok(columns) => { + for (tx, column) in column_streams_tx.iter().zip_eq(columns.into_iter()) { + if tx.send(Ok(column)).await.is_err() { + vortex_bail!( + "table column writer finished before all chunks were sent" + ); } } - Err(e) => { - let e: Arc = Arc::new(e); - for tx in column_streams_tx.iter() { - let _ = tx.send(Err(VortexError::from(Arc::clone(&e)))).await; - } - break; + } + Err(e) => { + let e: Arc = Arc::new(e); + for tx in column_streams_tx.iter() { + let _ = tx.send(Err(VortexError::from(Arc::clone(&e)))).await; } + return Err(VortexError::from(e)); } } - }) - .detach(); + } + Ok(()) + }; // First child column is the validity, subsequence children are the individual struct fields let column_dtypes: Vec = if is_nullable { @@ -359,7 +363,7 @@ impl LayoutStrategy for TableStrategy { }) .collect(); - let column_layouts = try_join_all(layout_futures).await?; + let (_success, column_layouts) = try_join(fanout_fut, try_join_all(layout_futures)).await?; // TODO(os): transposed stream could count row counts as well, // This must hold though, all columns must have the same row count of the struct layout let row_count = column_layouts.first().map(|l| l.row_count()).unwrap_or(0); @@ -370,12 +374,28 @@ impl LayoutStrategy for TableStrategy { #[cfg(test)] mod tests { use std::sync::Arc; + use std::task::Poll; + use vortex_array::ArrayContext; + use vortex_array::ArrayRef; + use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::StructFields; use vortex_array::field_path; + use vortex_error::VortexResult; + use vortex_io::runtime::single::block_on; + use vortex_io::session::RuntimeSessionExt; + use crate::LayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::table::TableStrategy; + use crate::segments::TestSegments; + use crate::sequence::SequenceId; + use crate::sequence::SequentialStreamAdapter; + use crate::sequence::SequentialStreamExt; + use crate::test::SESSION; #[test] #[should_panic( @@ -407,4 +427,41 @@ mod tests { ) .with_field_writer(FieldPath::root(), flat); } + + #[test] + #[should_panic(expected = "panic while transposing table stream")] + fn table_fanout_panic_propagates() { + let ctx = ArrayContext::empty(); + let segments = Arc::new(TestSegments::default()); + let (_, eof) = SequenceId::root().split(); + let dtype = DType::Struct( + StructFields::from_iter([( + "a", + DType::Primitive(PType::I32, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ); + let stream = + futures::stream::poll_fn(|_| -> Poll>> { + panic!("panic while transposing table stream"); + }); + let strategy = TableStrategy::new( + Arc::new(FlatLayoutStrategy::default()), + Arc::new(FlatLayoutStrategy::default()), + ); + + block_on(|handle| async move { + let session = SESSION.clone().with_handle(handle); + strategy + .write_stream( + ctx, + segments, + SequentialStreamAdapter::new(dtype, stream).sendable(), + eof, + &session, + ) + .await + .unwrap(); + }); + } }