diff --git a/rust/lance-file/src/writer.rs b/rust/lance-file/src/writer.rs index 9476c2fd29..e232530033 100644 --- a/rust/lance-file/src/writer.rs +++ b/rust/lance-file/src/writer.rs @@ -35,6 +35,7 @@ use lance_io::object_writer::ObjectWriter; use lance_io::traits::{WriteExt, Writer}; use object_store::path::Path; use snafu::{location, Location}; +use tokio::io::AsyncWriteExt; use crate::format::metadata::{Metadata, StatisticsMetadata}; use crate::format::{MAGIC, MAJOR_VERSION, MINOR_VERSION}; @@ -193,6 +194,12 @@ impl FileWriter { let batch_length = batches.iter().map(|b| b.num_rows() as i32).sum(); self.metadata.push_batch_length(batch_length); + // It's imperative we complete any in-flight requests, since we are + // returning control to the caller. If the caller takes a long time to + // write the next batch, the in-flight requests will not be polled and + // may time out. + self.object_writer.flush().await?; + self.batch_id += 1; Ok(()) }