feat: Add async-no-send feature flag to parquet

Author: jpopesculian

parquet/Cargo.toml

@@ -109,6 +109,8 @@ test_common = ["arrow/test_utils"]
 experimental = []
 # Enable async APIs
 async = ["futures", "tokio"]
+# Enable `!Send` with async APIs
+async-no-send = ["async"]
 # Enable object_store integration
 object_store = ["dep:object_store", "async"]
 # Group Zstd dependencies

parquet/src/arrow/async_reader/metadata.rs

@@ -21,9 +21,8 @@ use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
 use crate::file::page_index::index::Index;
 use crate::file::page_index::index_reader::{acc_range, decode_column_index, decode_offset_index};
 use crate::file::FOOTER_SIZE;
+use crate::util::async_util::{BoxFuture, Send};
 use bytes::Bytes;
-use futures::future::BoxFuture;
-use futures::FutureExt;
 use std::future::Future;
 use std::ops::Range;
 
@@ -66,7 +65,7 @@ pub trait MetadataFetch {
     /// Return a future that fetches the specified range of bytes asynchronously
     ///
     /// Note the returned type is a boxed future, often created by
-    /// [FutureExt::boxed]. See the trait documentation for an example
+    /// [futures::future::FutureExt::boxed]. See the trait documentation for an example
     fn fetch(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>>;
 }
 
@@ -82,7 +81,7 @@ pub trait MetadataSuffixFetch: MetadataFetch {
     /// Return a future that fetches the last `n` bytes asynchronously
     ///
     /// Note the returned type is a boxed future, often created by
-    /// [FutureExt::boxed]. See the trait documentation for an example
+    /// [futures::future::FutureExt::boxed]. See the trait documentation for an example
     fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result<Bytes>>;
 }
 
@@ -267,7 +266,7 @@ where
     Fut: Future<Output = Result<Bytes>> + Send,
 {
     fn fetch(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
-        async move { self.0(range.start.try_into()?..range.end.try_into()?).await }.boxed()
+        Box::pin(async move { self.0(range.start.try_into()?..range.end.try_into()?).await })
     }
 }
 

parquet/src/arrow/async_reader/mod.rs

@@ -30,7 +30,7 @@ use std::sync::Arc;
 use std::task::{Context, Poll};
 
 use bytes::{Buf, Bytes};
-use futures::future::{BoxFuture, FutureExt};
+use futures::future::FutureExt;
 use futures::ready;
 use futures::stream::Stream;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt};
@@ -54,6 +54,7 @@ use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
 use crate::file::page_index::offset_index::OffsetIndexMetaData;
 use crate::file::reader::{ChunkReader, Length, SerializedPageReader};
 use crate::format::{BloomFilterAlgorithm, BloomFilterCompression, BloomFilterHash};
+use crate::util::async_util::{BoxFuture, Send};
 
 mod metadata;
 pub use metadata::*;
@@ -85,7 +86,7 @@ pub trait AsyncFileReader: Send {
 
     /// Retrieve multiple byte ranges. The default implementation will call `get_bytes` sequentially
     fn get_byte_ranges(&mut self, ranges: Vec<Range<u64>>) -> BoxFuture<'_, Result<Vec<Bytes>>> {
-        async move {
+        Box::pin(async move {
             let mut result = Vec::with_capacity(ranges.len());
 
             for range in ranges.into_iter() {
@@ -94,8 +95,7 @@ pub trait AsyncFileReader: Send {
             }
 
             Ok(result)
-        }
-        .boxed()
+        })
     }
 
     /// Return a future which results in the [`ParquetMetaData`] for this Parquet file.
@@ -140,19 +140,18 @@ impl AsyncFileReader for Box<dyn AsyncFileReader + '_> {
 
 impl<T: AsyncFileReader + MetadataFetch + AsyncRead + AsyncSeek + Unpin> MetadataSuffixFetch for T {
     fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result<Bytes>> {
-        async move {
+        Box::pin(async move {
             self.seek(SeekFrom::End(-(suffix as i64))).await?;
             let mut buf = Vec::with_capacity(suffix);
             self.take(suffix as _).read_to_end(&mut buf).await?;
             Ok(buf.into())
-        }
-        .boxed()
+        })
     }
 }
 
 impl<T: AsyncRead + AsyncSeek + Unpin + Send> AsyncFileReader for T {
     fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
-        async move {
+        Box::pin(async move {
             self.seek(SeekFrom::Start(range.start)).await?;
 
             let to_read = range.end - range.start;
@@ -163,15 +162,14 @@ impl<T: AsyncRead + AsyncSeek + Unpin + Send> AsyncFileReader for T {
             }
 
             Ok(buffer.into())
-        }
-        .boxed()
+        })
     }
 
     fn get_metadata<'a>(
         &'a mut self,
         options: Option<&'a ArrowReaderOptions>,
     ) -> BoxFuture<'a, Result<Arc<ParquetMetaData>>> {
-        async move {
+        Box::pin(async move {
             let metadata_reader = ParquetMetaDataReader::new()
                 .with_page_indexes(options.is_some_and(|o| o.page_index));
 
@@ -182,8 +180,7 @@ impl<T: AsyncRead + AsyncSeek + Unpin + Send> AsyncFileReader for T {
 
             let parquet_metadata = metadata_reader.load_via_suffix_and_finish(self).await?;
             Ok(Arc::new(parquet_metadata))
-        }
-        .boxed()
+        })
     }
 }
 
@@ -844,14 +841,12 @@ where
 
                     let selection = self.selection.as_mut().map(|s| s.split_off(row_count));
 
-                    let fut = reader
-                        .read_row_group(
-                            row_group_idx,
-                            selection,
-                            self.projection.clone(),
-                            self.batch_size,
-                        )
-                        .boxed();
+                    let fut = Box::pin(reader.read_row_group(
+                        row_group_idx,
+                        selection,
+                        self.projection.clone(),
+                        self.batch_size,
+                    ));
 
                     self.state = StreamState::Reading(fut)
                 }

parquet/src/arrow/async_reader/store.rs

@@ -21,8 +21,9 @@ use crate::arrow::arrow_reader::ArrowReaderOptions;
 use crate::arrow::async_reader::{AsyncFileReader, MetadataSuffixFetch};
 use crate::errors::{ParquetError, Result};
 use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
+use crate::util::async_util::BoxFuture;
 use bytes::Bytes;
-use futures::{future::BoxFuture, FutureExt, TryFutureExt};
+use futures::{FutureExt, TryFutureExt};
 use object_store::{path::Path, ObjectStore};
 use object_store::{GetOptions, GetRange};
 use tokio::runtime::Handle;
@@ -59,6 +60,7 @@ pub struct ParquetObjectReader {
     metadata_size_hint: Option<usize>,
     preload_column_index: bool,
     preload_offset_index: bool,
+    #[cfg_attr(feature = "async-no-send", allow(dead_code))]
     runtime: Option<Handle>,
 }
 
@@ -132,6 +134,18 @@ impl ParquetObjectReader {
         }
     }
 
+    #[cfg(feature = "async-no-send")]
+    fn spawn<F, O, E>(&self, f: F) -> BoxFuture<'_, Result<O>>
+    where
+        F: for<'a> FnOnce(&'a Arc<dyn ObjectStore>, &'a Path) -> BoxFuture<'a, Result<O, E>>
+            + 'static,
+        O: 'static,
+        E: Into<ParquetError> + 'static,
+    {
+        Box::pin(f(&self.store, &self.path).map_err(|e| e.into()))
+    }
+
+    #[cfg(not(feature = "async-no-send"))]
     fn spawn<F, O, E>(&self, f: F) -> BoxFuture<'_, Result<O>>
     where
         F: for<'a> FnOnce(&'a Arc<dyn ObjectStore>, &'a Path) -> BoxFuture<'a, Result<O, E>>
@@ -144,18 +158,19 @@ impl ParquetObjectReader {
             Some(handle) => {
                 let path = self.path.clone();
                 let store = Arc::clone(&self.store);
-                handle
-                    .spawn(async move { f(&store, &path).await })
-                    .map_ok_or_else(
-                        |e| match e.try_into_panic() {
-                            Err(e) => Err(ParquetError::External(Box::new(e))),
-                            Ok(p) => std::panic::resume_unwind(p),
-                        },
-                        |res| res.map_err(|e| e.into()),
-                    )
-                    .boxed()
+                Box::pin(
+                    handle
+                        .spawn(async move { f(&store, &path).await })
+                        .map_ok_or_else(
+                            |e| match e.try_into_panic() {
+                                Err(e) => Err(ParquetError::External(Box::new(e))),
+                                Ok(p) => std::panic::resume_unwind(p),
+                            },
+                            |res| res.map_err(|e| e.into()),
+                        ),
+                )
             }
-            None => f(&self.store, &self.path).map_err(|e| e.into()).boxed(),
+            None => Box::pin(f(&self.store, &self.path).map_err(|e| e.into())),
         }
     }
 }
@@ -181,10 +196,7 @@ impl AsyncFileReader for ParquetObjectReader {
         self.spawn(|store, path| store.get_range(path, range))
     }
 
-    fn get_byte_ranges(&mut self, ranges: Vec<Range<u64>>) -> BoxFuture<'_, Result<Vec<Bytes>>>
-    where
-        Self: Send,
-    {
+    fn get_byte_ranges(&mut self, ranges: Vec<Range<u64>>) -> BoxFuture<'_, Result<Vec<Bytes>>> {
         self.spawn(|store, path| async move { store.get_ranges(path, &ranges).await }.boxed())
     }
 

parquet/src/arrow/async_writer/mod.rs

@@ -66,12 +66,11 @@ use crate::{
     errors::{ParquetError, Result},
     file::{metadata::RowGroupMetaData, properties::WriterProperties},
     format::{FileMetaData, KeyValue},
+    util::async_util::{BoxFuture, Send},
 };
 use arrow_array::RecordBatch;
 use arrow_schema::SchemaRef;
 use bytes::Bytes;
-use futures::future::BoxFuture;
-use futures::FutureExt;
 use std::mem;
 use tokio::io::{AsyncWrite, AsyncWriteExt};
 
@@ -103,20 +102,18 @@ impl AsyncFileWriter for Box<dyn AsyncFileWriter + '_> {
 
 impl<T: AsyncWrite + Unpin + Send> AsyncFileWriter for T {
     fn write(&mut self, bs: Bytes) -> BoxFuture<'_, Result<()>> {
-        async move {
+        Box::pin(async move {
             self.write_all(&bs).await?;
             Ok(())
-        }
-        .boxed()
+        })
     }
 
     fn complete(&mut self) -> BoxFuture<'_, Result<()>> {
-        async move {
+        Box::pin(async move {
             self.flush().await?;
             self.shutdown().await?;
             Ok(())
-        }
-        .boxed()
+        })
     }
 }
 

parquet/src/arrow/async_writer/store.rs

@@ -16,11 +16,11 @@
 // under the License.
 
 use bytes::Bytes;
-use futures::future::BoxFuture;
 use std::sync::Arc;
 
 use crate::arrow::async_writer::AsyncFileWriter;
 use crate::errors::{ParquetError, Result};
+use crate::util::async_util::BoxFuture;
 use object_store::buffered::BufWriter;
 use object_store::path::Path;
 use object_store::ObjectStore;

parquet/src/util/async_util.rs

@@ -0,0 +1,15 @@
+#[cfg(feature = "async-no-send")]
+mod send_impl {
+    pub trait Send {}
+    impl<T> Send for T {}
+    pub type BoxFuture<'a, T> = futures::future::LocalBoxFuture<'a, T>;
+}
+
+#[cfg(not(feature = "async-no-send"))]
+mod send_impl {
+    pub trait Send: std::marker::Send {}
+    impl<T> Send for T where T: std::marker::Send {}
+    pub type BoxFuture<'a, T> = futures::future::BoxFuture<'a, T>;
+}
+
+pub use send_impl::*;

parquet/src/util/mod.rs

@@ -24,6 +24,9 @@ pub(crate) mod interner;
 pub(crate) mod test_common;
 pub mod utf8;
 
+#[cfg(feature = "async")]
+pub(crate) mod async_util;
+
 #[cfg(any(test, feature = "test_common"))]
 pub use self::test_common::page_util::{
     DataPageBuilder, DataPageBuilderImpl, InMemoryPageIterator,