feat(rust): ParquetCloudSink to allow streaming pipelines into remote ObjectStores

crates/Makefile

@@ -118,6 +118,7 @@ check-wasm:  ## Check wasm build without supported features
 		--exclude-features async              \
 		--exclude-features aws                \
 		--exclude-features azure              \
+    --exclude-features cloud              \
 		--exclude-features decompress         \
 		--exclude-features decompress-fast    \
 		--exclude-features default            \

crates/polars-core/Cargo.toml

@@ -181,7 +181,7 @@ docs-selection = [
 ]
 
 # Cloud support.
-"async" = ["url"]
+"async" = ["url", "object_store"]
 "aws" = ["async", "object_store/aws"]
 "azure" = ["async", "object_store/azure"]
 "gcp" = ["async", "object_store/gcp"]

crates/polars-error/Cargo.toml

@@ -10,5 +10,6 @@ description = "Error definitions for the Polars DataFrame library"
 
 [dependencies]
 arrow = { workspace = true }
+object_store = { version = "0.6.0", default-features = false, optional = true }
 regex = { workspace = true, optional = true }
 thiserror = { workspace = true }

crates/polars-error/src/lib.rs

@@ -79,6 +79,16 @@ impl From<regex::Error> for PolarsError {
     }
 }
 
+#[cfg(feature = "object_store")]
+impl From<object_store::Error> for PolarsError {
+    fn from(err: object_store::Error) -> Self {
+        PolarsError::Io(std::io::Error::new(
+            std::io::ErrorKind::Other,
+            format!("object store error {err:?}"),
+        ))
+    }
+}
+
 pub type PolarsResult<T> = Result<T, PolarsError>;
 
 pub use arrow::error::Error as ArrowError;

crates/polars-io/Cargo.toml

@@ -39,6 +39,7 @@ serde_json = { version = "1", default-features = false, features = ["alloc", "ra
 simd-json = { workspace = true, optional = true }
 simdutf8 = { version = "0.1", optional = true }
 tokio = { version = "1.26", features = ["net"], optional = true }
+tokio-util = { version = "0.7.8", features = ["io", "io-util"], optional = true }
 url = { workspace = true, optional = true }
 
 [target.'cfg(not(target_family = "wasm"))'.dependencies]
@@ -87,8 +88,8 @@ dtype-decimal = ["polars-core/dtype-decimal"]
 fmt = ["polars-core/fmt"]
 lazy = []
 parquet = ["polars-core/parquet", "arrow/io_parquet", "arrow/io_parquet_compression", "memmap"]
-async = ["async-trait", "futures", "tokio", "arrow/io_ipc_write_async", "polars-error/regex"]
-cloud = ["object_store", "async", "url"]
+async = ["async-trait", "futures", "tokio", "tokio-util", "arrow/io_ipc_write_async", "polars-error/regex"]
+cloud = ["object_store", "async", "polars-core/async", "polars-error/object_store", "url"]
 aws = ["object_store/aws", "cloud", "polars-core/aws"]
 azure = ["object_store/azure", "cloud", "polars-core/azure"]
 gcp = ["object_store/gcp", "cloud", "polars-core/gcp"]

crates/polars-io/src/cloud/adaptors.rs

@@ -10,7 +10,10 @@ use futures::future::BoxFuture;
 use futures::lock::Mutex;
 use futures::{AsyncRead, AsyncSeek, Future, TryFutureExt};
 use object_store::path::Path;
-use object_store::ObjectStore;
+use object_store::{MultipartId, ObjectStore};
+use polars_core::cloud::CloudOptions;
+use polars_error::{PolarsError, PolarsResult};
+use tokio::io::{AsyncWrite, AsyncWriteExt};
 
 type OptionalFuture = Arc<Mutex<Option<BoxFuture<'static, std::io::Result<Vec<u8>>>>>>;
 
@@ -129,3 +132,159 @@ impl AsyncSeek for CloudReader {
         std::task::Poll::Ready(Ok(self.pos))
     }
 }
+
+/// Adaptor which wraps the asynchronous interface of [ObjectStore::put_multipart](https://docs.rs/object_store/latest/object_store/trait.ObjectStore.html#tymethod.put_multipart)
+/// exposing a synchronous interface which implements `std::io::Write`.
+///
+/// This allows it to be used in sync code which would otherwise write to a simple File or byte stream,
+/// such as with `polars::prelude::CsvWriter`.
+pub struct CloudWriter {
+    // Hold a reference to the store
+    object_store: Arc<dyn ObjectStore>,
+    // The path in the object_store which we want to write to
+    path: Path,
+    // ID of a partially-done upload, used to abort the upload on error
+    multipart_id: MultipartId,
+    // The Tokio runtime which the writer uses internally.
+    runtime: tokio::runtime::Runtime,
+    // Internal writer, constructed at creation
+    writer: Box<dyn AsyncWrite + Send + Unpin>,
+}
+
+impl CloudWriter {
+    /// Construct a new CloudWriter, re-using the given `object_store`
+    ///
+    /// Creates a new (current-thread) Tokio runtime
+    /// which bridges the sync writing process with the async ObjectStore multipart uploading.
+    /// TODO: Naming?
+    pub fn new_with_object_store(
+        object_store: Arc<dyn ObjectStore>,
+        path: Path,
+    ) -> PolarsResult<Self> {
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .enable_io()
+            .enable_time()
+            .build()
+            .unwrap();
+        let build_result =
+            runtime.block_on(async { Self::build_writer(&object_store, &path).await });
+        match build_result {
+            Err(error) => Err(PolarsError::from(error)),
+            Ok((multipart_id, writer)) => Ok(CloudWriter {
+                object_store,
+                path,
+                multipart_id,
+                runtime,
+                writer,
+            }),
+        }
+    }
+
+    /// Constructs a new CloudWriter from a path and an optional set of CloudOptions.
+    ///
+    /// Wrapper around `CloudWriter::new_with_object_store` that is useful if you only have a single write task.
+    /// TODO: Naming?
+    pub fn new(uri: &str, cloud_options: Option<&CloudOptions>) -> PolarsResult<Self> {
+        let (cloud_location, object_store) = crate::cloud::build(uri, cloud_options)?;
+        let object_store = Arc::from(object_store);
+        Self::new_with_object_store(object_store, cloud_location.prefix.into())
+    }
+
+    async fn build_writer(
+        object_store: &Arc<dyn ObjectStore>,
+        path: &Path,
+    ) -> object_store::Result<(MultipartId, Box<dyn AsyncWrite + Send + Unpin>)> {
+        let (multipart_id, s3_writer) = object_store.put_multipart(path).await?;
+        Ok((multipart_id, s3_writer))
+    }
+
+    fn abort(&self) {
+        let _ = self.runtime.block_on(async {
+            self.object_store
+                .abort_multipart(&self.path, &self.multipart_id)
+                .await
+        });
+    }
+}
+
+impl std::io::Write for CloudWriter {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        let res = self.runtime.block_on(self.writer.write(buf));
+        if res.is_err() {
+            self.abort();
+        }
+        res
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        let res = self.runtime.block_on(self.writer.flush());
+        if res.is_err() {
+            self.abort();
+        }
+        res
+    }
+}
+
+impl Drop for CloudWriter {
+    fn drop(&mut self) {
+        let _ = self.runtime.block_on(self.writer.shutdown());
+    }
+}
+
+#[cfg(feature = "csv")]
+#[cfg(test)]
+mod tests {
+    use object_store::ObjectStore;
+    use polars_core::df;
+    use polars_core::prelude::{DataFrame, NamedFrom};
+
+    use super::*;
+
+    fn example_dataframe() -> DataFrame {
+        df!(
+            "foo" => &[1, 2, 3],
+            "bar" => &[None, Some("bak"), Some("baz")],
+        )
+        .unwrap()
+    }
+
+    // Skip this tests on Windows since it does not have a convenient /tmp/ location.
+    #[cfg_attr(target_os = "windows", ignore)]
+    #[test]
+    fn csv_to_local_objectstore_cloudwriter() {
+        use crate::csv::CsvWriter;
+        use crate::prelude::SerWriter;
+
+        let mut df = example_dataframe();
+
+        let object_store: Box<dyn ObjectStore> = Box::new(
+            object_store::local::LocalFileSystem::new_with_prefix("/tmp/")
+                .expect("Could not initialize connection"),
+        );
+        let object_store: Arc<dyn ObjectStore> = Arc::from(object_store);
+
+        let path: object_store::path::Path = "cloud_writer_example.csv".into();
+
+        let mut cloud_writer = CloudWriter::new_with_object_store(object_store, path).unwrap();
+        CsvWriter::new(&mut cloud_writer)
+            .finish(&mut df)
+            .expect("Could not write dataframe as CSV to remote location");
+    }
+
+    // Skip this tests on Windows since it does not have a convenient /tmp/ location.
+    #[cfg_attr(target_os = "windows", ignore)]
+    #[test]
+    fn cloudwriter_from_cloudlocation_test() {
+        use crate::csv::CsvWriter;
+        use crate::prelude::SerWriter;
+
+        let mut df = example_dataframe();
+
+        let mut cloud_writer =
+            CloudWriter::new("file:///tmp/cloud_writer_example2.csv", None).unwrap();
+
+        CsvWriter::new(&mut cloud_writer)
+            .finish(&mut df)
+            .expect("Could not write dataframe as CSV to remote location");
+    }
+}

crates/polars-io/src/lib.rs

@@ -5,7 +5,7 @@
 #[cfg(feature = "avro")]
 pub mod avro;
 #[cfg(feature = "cloud")]
-mod cloud;
+pub mod cloud;
 #[cfg(any(feature = "csv", feature = "json"))]
 pub mod csv;
 #[cfg(feature = "parquet")]

crates/polars-lazy/Cargo.toml

@@ -45,6 +45,7 @@ async = [
   "polars-pipe/async",
   "streaming",
 ]
+cloud = ["async", "polars-pipe/cloud"]
 ipc = ["polars-io/ipc", "polars-plan/ipc", "polars-pipe/ipc"]
 json = ["polars-io/json", "polars-plan/json", "polars-json"]
 csv = ["polars-io/csv", "polars-plan/csv", "polars-pipe/csv"]

crates/polars-lazy/src/frame/mod.rs

@@ -579,9 +579,9 @@ impl LazyFrame {
             None
         };
 
-        // file sink should be replaced
+        // sink should be replaced
         let no_file_sink = if check_sink {
-            !matches!(lp_arena.get(lp_top), ALogicalPlan::FileSink { .. })
+            !matches!(lp_arena.get(lp_top), ALogicalPlan::Sink { .. })
         } else {
             true
         };
@@ -639,9 +639,9 @@ impl LazyFrame {
     #[cfg(feature = "parquet")]
     pub fn sink_parquet(mut self, path: PathBuf, options: ParquetWriteOptions) -> PolarsResult<()> {
         self.opt_state.streaming = true;
-        self.logical_plan = LogicalPlan::FileSink {
+        self.logical_plan = LogicalPlan::Sink {
             input: Box::new(self.logical_plan),
-            payload: FileSinkOptions {
+            payload: SinkType::File {
                 path: Arc::new(path),
                 file_type: FileType::Parquet(options),
             },
@@ -656,15 +656,44 @@ impl LazyFrame {
         Ok(())
     }
 
+    /// Stream a query result into a parquet file on an ObjectStore-compatible cloud service. This is useful if the final result doesn't fit
+    /// into memory, and where you do not want to write to a local file but to a location in the cloud.
+    /// This method will return an error if the query cannot be completely done in a
+    /// streaming fashion.
+    #[cfg(all(feature = "cloud", feature = "parquet"))]
+    pub fn sink_parquet_cloud(
+        mut self,
+        uri: String,
+        cloud_options: Option<polars_core::cloud::CloudOptions>,
+        parquet_options: ParquetWriteOptions,
+    ) -> PolarsResult<()> {
+        self.opt_state.streaming = true;
+        self.logical_plan = LogicalPlan::Sink {
+            input: Box::new(self.logical_plan),
+            payload: SinkType::Cloud {
+                uri: Arc::new(uri),
+                cloud_options,
+                file_type: FileType::Parquet(parquet_options),
+            },
+        };
+        let (mut state, mut physical_plan, is_streaming) = self.prepare_collect(true)?;
+        polars_ensure!(
+            is_streaming,
+            ComputeError: "cannot run the whole query in a streaming order"
+        );
+        let _ = physical_plan.execute(&mut state)?;
+        Ok(())
+    }
+
     /// Stream a query result into an ipc/arrow file. This is useful if the final result doesn't fit
     /// into memory. This methods will return an error if the query cannot be completely done in a
     /// streaming fashion.
     #[cfg(feature = "ipc")]
     pub fn sink_ipc(mut self, path: PathBuf, options: IpcWriterOptions) -> PolarsResult<()> {
         self.opt_state.streaming = true;
-        self.logical_plan = LogicalPlan::FileSink {
+        self.logical_plan = LogicalPlan::Sink {
             input: Box::new(self.logical_plan),
-            payload: FileSinkOptions {
+            payload: SinkType::File {
                 path: Arc::new(path),
                 file_type: FileType::Ipc(options),
             },

crates/polars-lazy/src/physical_plan/planner/lp.rs

@@ -150,9 +150,16 @@ pub fn create_physical_plan(
     match logical_plan {
         #[cfg(feature = "python")]
         PythonScan { options, .. } => Ok(Box::new(executors::PythonScanExec { options })),
-        FileSink { .. } => panic!(
-            "sink_parquet not yet supported in standard engine. Use 'collect().write_parquet()'"
-        ),
+        Sink { payload, .. } => {
+            match payload {
+                SinkType::Memory => panic!("Memory Sink not supported in the standard engine."),
+                SinkType::File{..} => panic!(
+                    "sink_parquet not yet supported in standard engine. Use 'collect().write_parquet()'"
+                ),
+                #[cfg(feature = "cloud")]
+                SinkType::Cloud{..} => panic!("Cloud Sink not supported in standard engine.")
+            }
+        }
         Union { inputs, options } => {
             let inputs = inputs
                 .into_iter()

crates/polars-lazy/src/physical_plan/streaming/construct_pipeline.rs

@@ -122,7 +122,7 @@ pub(super) fn construct(
         // the file sink is always to the top of the tree
         // not every branch has a final sink. For instance rhs join branches
         if let Some(node) = branch.get_final_sink() {
-            if matches!(lp_arena.get(node), ALogicalPlan::FileSink { .. }) {
+            if matches!(lp_arena.get(node), ALogicalPlan::Sink { .. }) {
                 final_sink = Some(node)
             }
         }
@@ -191,20 +191,16 @@ pub(super) fn construct(
         return Ok(None);
     };
     let insertion_location = match lp_arena.get(final_sink) {
-        FileSink {
+        // this was inserted only during conversion and does not exist
+        // in the original tree, so we take the input, as that's where
+        // we connect into the original tree.
+        Sink {
             input,
-            payload: FileSinkOptions { file_type, .. },
-        } => {
-            // this was inserted only during conversion and does not exist
-            // in the original tree, so we take the input, as that's where
-            // we connect into the original tree.
-            if matches!(file_type, FileType::Memory) {
-                *input
-            } else {
-                // default case if the tree ended with a file_sink
-                final_sink
-            }
-        },
+            payload: SinkType::Memory,
+        } => *input,
+        // Other sinks were not inserted during conversion,
+        // so they are returned as-is
+        Sink { .. } => final_sink,
         _ => unreachable!(),
     };
     // keep the original around for formatting purposes